Skip to content

[PyTorch] Remove is_first_microbatch setting after cudagraph warmup #2715

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yaox12 merged 2 commits into from
Mar 2, 2026
Merged

[PyTorch] Remove is_first_microbatch setting after cudagraph warmup #2715

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
25af05a
Remove is_first_microbatch setting after warmup
buptzyb Feb 27, 2026
7abebf0
Merge branch 'main' into robinz/graph_isfirstmicrobatch
yaox12 Mar 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 0 additions & 5 deletions transformer_engine/pytorch/graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -507,11 +507,6 @@ def hook_fn(
else:
grad_inputs = None
del outputs, grad_inputs
# The following code is added specifically for MCore's special requirements,
# aimed at preventing warmup from altering the control flow.
for module in func.modules():
if hasattr(module, "is_first_microbatch"):
module.is_first_microbatch = True
torch.cuda.synchronize()
Copy link
Contributor

@greptile-apps greptile-apps bot Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


# All captures here share a mempool. To avoid replays corrupting each other's memory,
Expand Down
Loading