Skip to content

fix(gtp): clear dummy gradients before DDP readiness - #7467

Merged
ericharper merged 1 commit into
NVIDIA:mainfrom
JF-D:fix/gtp-clear-dummy-before-ddp
Oct 10, 2026
Merged

ericharper merged 1 commit into
NVIDIA:mainfrom
JF-D:fix/gtp-clear-dummy-before-ddp

Conversation

@JF-D

@JF-D JF-D commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Clear GTP's autograd placeholder before invoking its manual DDP readiness callback.

GTP has already accumulated the reduced gradient into main_grad. However, when
zero_out_wgrad=True (for example, on an MTP embedding), DDP can still add an old
param.grad even though grad_added_to_main_grad=True. Under CUDA graph capture,
that redundant addition can record the address of a warmup allocation that is
released before replay.

Set param.grad = None at the GTP notification boundary. Preserve the
zero_out_wgrad convention, the placeholder returned to autograd, capture
bookkeeping, readiness notifications, and generic DDP accumulation behavior.

Issue tracking

Linked issue: None; regression reproduced by the included tests.

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
  • I have added proper typing to my code Typing guidelines
  • I have added relevant documentation (callback contract and regression-test docstrings)
  • I have run the autoformatter.sh on my PR

Code review

The production change is confined to GTP. No GPU kernels or public APIs are added.

@copy-pr-bot

copy-pr-bot Bot commented Sep 18, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@JF-D
JF-D marked this pull request as ready for review September 18, 2026 05:07
@JF-D
JF-D requested review from a team as code owners September 18, 2026 05:07
@JF-D JF-D self-assigned this Sep 18, 2026
@JF-D
JF-D force-pushed the fix/gtp-clear-dummy-before-ddp branch from af3cf8f to 2c48c5f Compare September 18, 2026 06:28
@JF-D JF-D changed the title fix(gtp): clear dummy gradients before DDP readiness fix(gtp): skip dummy accumulation for shared weights Sep 18, 2026
GTP has already accumulated its reduced gradient into main_grad when it
invokes the manual DDP callback. A stale autograd placeholder can still
be added when zero_out_wgrad overrides grad_added_to_main_grad. Capturing
that redundant add can retain a warmup allocation after it is released.

Clear the placeholder at the GTP notification boundary while preserving
the dummy returned to autograd and generic DDP accumulation semantics.
Add regression coverage for ordinary DDP accumulation, repeated GTP
backward uses, and replay after releasing the warmup dummy.

Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
@JF-D
JF-D force-pushed the fix/gtp-clear-dummy-before-ddp branch from 2c48c5f to 2b8404f Compare September 18, 2026 06:41
@JF-D JF-D changed the title fix(gtp): skip dummy accumulation for shared weights fix(gtp): clear dummy gradients before DDP readiness Sep 18, 2026
@JF-D

JF-D commented Oct 1, 2026

Copy link
Copy Markdown
Contributor Author

/review

@svcnvidia-nemo-ci svcnvidia-nemo-ci left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM. The GTP-specific change preserves DDP readiness notifications and ordinary gradient accumulation, with regression coverage for stale placeholders, repeated backward use, and CUDA graph replay. Tests were not executed in this read-only review.

Review generated by codex

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Oct 9, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Oct 9, 2026
@ericharper
ericharper enabled auto-merge October 9, 2026 21:57
@ericharper
ericharper added this pull request to the merge queue Oct 9, 2026
@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 6cad8f6 Oct 10, 2026
114 of 124 checks passed
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Oct 14, 2026
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit 6cad8f6)
Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Oct 14, 2026
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit 6cad8f6)
Signed-off-by: xuwenc <xuwenc@nvidia.com>
@JF-D
JF-D deleted the fix/gtp-clear-dummy-before-ddp branch October 16, 2026 22:16
fanshiqing pushed a commit to fanshiqing/Megatron-LM that referenced this pull request Oct 27, 2026
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>

Upstream: NVIDIA#7467

Adapt the manual-hook call context to this support branch; retain its existing
captured-gradient readiness implementation. Include the upstream regression tests unchanged.

(cherry picked from commit 6cad8f6)
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

This branch was previously deployed

2 inactive deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: low

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants