remove dead manual_release_grads code path in 1F1B overlap schedule by Wohox · Pull Request #4511 · NVIDIA/Megatron-LM

Wohox · 2026-04-29T01:51:26Z

What does this PR do ?

Removes dead code in the 1F1B overlap schedule. ScheduleNode.manual_release_grads was hardcoded to False at construction in pipeline_parallel/utils.py and never reassigned anywhere in the codebase, so both if self.manual_release_grads: branches were unreachable. The supporting bookkeeping (delay_grads_release, output_grads) existed only to feed those unreachable branches.

Removed:

pipeline_parallel/utils.py: self.manual_release_grads = False and self.delay_grads_release = False fields on ScheduleNode; the unreachable g.untyped_storage().resize_(0) branch in _backward.
models/gpt/fine_grained_callables.py: in backward_impl, the if self.delay_wgrad_compute: self.output_grads = grads; self.delay_grads_release = ... block (only consumed by the dead release path); in backward_dw, the unreachable tensor.untyped_storage().resize_(0) loop, the now-meaningless assert self.delay_grads_release, and the explicit self.output_grads = None clearing.

PyTorch's caching allocator already reclaims the dgrad tensors when the last Python reference drops at the end of _backward / backward_dw, so this whole machinery was a no-op. Same scope as #4122.

Why this is safe

A/B comparison was run with the flag temporarily exposed as a TransformerConfig field, on qwen3_moe_proxy (PP=1 EP=8, overlap_moe_expert_parallel_comm + delay_wgrad_compute enabled, 100 iters on real 8×H100, EOS):

`manual_release_grads`	Peak max-allocated (MB)
`False` (status quo)	12479.13
`True` (release on)	12479.13

Peak allocated memory is identical to two decimals.

Companion dev PR: #4510

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

(Pure deletion of unreachable code; existing 1F1B-overlap functional tests cover the surrounding paths.)

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

copy-pr-bot · 2026-04-29T01:51:30Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

`ScheduleNode.manual_release_grads` was hardcoded to False at construction and never reassigned anywhere in the codebase, so the two `if self.manual_release_grads:` branches in `ScheduleNode._backward` and `TransformerLayerNode.backward_dw` were unreachable. The supporting bookkeeping — `self.delay_grads_release`, `self.output_grads`, the `assert self.delay_grads_release` in `backward_dw`, and the `if self.delay_wgrad_compute: self.output_grads = grads; ...` block in `backward_impl` — existed only to feed the unreachable branches. PyTorch's caching allocator already reclaims the dgrad tensors when the last Python reference drops at the end of `_backward` / `backward_dw`, so this whole machinery was a no-op. A/B comparison with the flag temporarily exposed as a TransformerConfig field (qwen3_moe_proxy, PP=1 EP=8, overlap_moe_expert_parallel_comm + delay_wgrad_compute, 100 iters on real 8xH100): peak max-allocated memory was identical at 12479.13 MB with the flag both ON and OFF.

asolergi-nv · 2026-05-14T07:53:15Z

/ok to test 876f049

svcnvidia-nemo-ci · 2026-05-14T09:36:37Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25852953881

…apping_main Conflict resolution + porting upstream deltas to relocated code: - gpt/fine_grained_callables.py: keep slim HEAD; PreProcess / PostProcess / TransformerLayerNode / _BackwardDWWrapper and build_mtp_layer_callables / build_layer_callables now live under common/. - common/utils.py: port PR NVIDIA#4511 (remove dead manual_release_grads code path) into TransformerLayerNode.backward_impl / backward_dw. - common/fine_grained_callables.py: port PR NVIDIA#2645 (packed sequence GDN) into build_mtp_layer_callables — unpack the new 5-tuple from _get_embeddings and forward packed_seq_params / padding_mask. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>

Wohox force-pushed the denliu/remove-dead-manual-release-grads-main branch from 7d20639 to 065119b Compare April 29, 2026 02:00

Wohox mentioned this pull request Apr 29, 2026

[Dev] remove dead manual_release_grads code path in 1F1B overlap schedule #4510

Merged

5 tasks

Merge branch 'main' into denliu/remove-dead-manual-release-grads-main

6e75e79

Wohox marked this pull request as ready for review April 29, 2026 02:06

Wohox requested review from a team as code owners April 29, 2026 02:06

svcnvidia-nemo-ci requested a review from a team April 29, 2026 02:06

svcnvidia-nemo-ci added the complexity: low label Apr 29, 2026

yashaswikarnati approved these changes Apr 29, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Apr 29, 2026

ericharper approved these changes May 12, 2026

View reviewed changes

jaredcasper approved these changes May 14, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels May 14, 2026

Merge branch 'main' into denliu/remove-dead-manual-release-grads-main

876f049

asolergi-nv enabled auto-merge May 14, 2026 07:52

copy-pr-bot Bot temporarily deployed to public May 14, 2026 07:53 Inactive

copy-pr-bot Bot temporarily deployed to test May 14, 2026 07:54 Inactive

asolergi-nv added this pull request to the merge queue May 14, 2026

Merged via the queue into NVIDIA:main with commit 58d45bf May 14, 2026
68 of 70 checks passed

cspades pushed a commit to cspades/Megatron-LM that referenced this pull request May 14, 2026

remove dead manual_release_grads code path in 1F1B overlap schedule (N…

af6dc38

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

janEbert pushed a commit to janEbert/Megatron-LM that referenced this pull request Jun 2, 2026

remove dead manual_release_grads code path in 1F1B overlap schedule (N…

2762199

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

remove dead manual_release_grads code path in 1F1B overlap schedule (N…

96c1ee1

…VIDIA#4511) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Navigation Menu

Uh oh!

remove dead manual_release_grads code path in 1F1B overlap schedule - #4511

remove dead manual_release_grads code path in 1F1B overlap schedule#4511
asolergi-nv merged 3 commits into
NVIDIA:mainfrom
Wohox:denliu/remove-dead-manual-release-grads-main

Wohox commented Apr 29, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Apr 29, 2026

Uh oh!

asolergi-nv commented May 14, 2026

Uh oh!

svcnvidia-nemo-ci commented May 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

Wohox commented Apr 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Why this is safe

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Apr 29, 2026

Uh oh!

asolergi-nv commented May 14, 2026

Uh oh!

svcnvidia-nemo-ci commented May 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Wohox commented Apr 29, 2026 •

edited

Loading