[PyTorch] Remove is_first_microbatch setting after cudagraph warmup

[buptzyb]

Description

Reset is_first_microbatch to True is unnecessary for fp8, because later we'll set it back to False: https://github.com/NVIDIA/TransformerEngine/blob/release_v2.12/transformer_engine/pytorch/module/layernorm_mlp.py#L2037-L2042. It was used to control fp8 cast-transpose behavior, which is overtaken by the skip_fp8_weight_update tensor in graph replay.

For bf16, it's not only unnecessary but also conflicts with an optimization in mcore: when PP is not enabled, we capture only one graph per layer and share it across microbatches. is_first_microbatch controls the behavior of fused wgrad: D=A*B for the first microbatch, and D=A*B+D for the following microbatches. We should always capture the D=A*B+D version (i.e., capture with is_first_microbatch=False) to make the graph work for all microbatches. Without reset, it must be is_first_microbatch=False after warmup.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

Removes the post-warmup reset of is_first_microbatch attribute to True on modules. This code was originally added for Megatron-Core (MCore) compatibility to prevent warmup from altering control flow. The author argues this reset is unnecessary because during CUDA graph capture, the skip_fp8_weight_update tensor now controls FP8 weight update behavior, and is_first_microbatch is automatically set to False when this tensor is present. For wgrad accumulation, capturing the no-clear version (is_first_microbatch=False) allows the graph to work correctly across all microbatches rather than just the first one.

Confidence Score: 3/5

• Safe to merge with thorough testing, especially with MCore integration
• The change is small and well-explained, but it removes code specifically added for Megatron-Core compatibility without including tests to validate the new behavior. The reasoning about skip_fp8_weight_update controlling the behavior is sound, but the removed code operated on module attributes while the explanation focuses on function parameters. This needs validation with MCore workloads to ensure no regression.
• No files require special attention beyond standard MCore integration testing

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Removes post-warmup is_first_microbatch attribute reset that was added for MCore compatibility

_{Last reviewed commit: 7abebf0}

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[yaox12]

/te-ci pytorch

[greptile-apps]

Check that MCore modules no longer rely on the is_first_microbatch attribute being reset after warmup. The removed code handled module attributes while the PR description discusses function parameters - these are different mechanisms. Verify that removing this reset doesn't break MCore integration where modules may have is_first_microbatch as an instance attribute that affects control flow.

[yaox12]

LGTM.