Skip to content

fix(MTP's memory): Reuse the GTP wgrad send buffer across a weight's repeated backwards - #7470

Merged
fanshiqing merged 1 commit into
NVIDIA:mainfrom
fanshiqing:memory_opt/reuse_gpt_wgrad_send_buf_across_repeated_bwds
Sep 22, 2026
Merged

fix(MTP's memory): Reuse the GTP wgrad send buffer across a weight's repeated backwards#7470
fanshiqing merged 1 commit into
NVIDIA:mainfrom
fanshiqing:memory_opt/reuse_gpt_wgrad_send_buf_across_repeated_bwds

Conversation

@fanshiqing

@fanshiqing fanshiqing commented Sep 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

TL;DR

This MR help reduce the MTP's output_layer's GTP_reduce_scatter input buffer memory cost:

  • vocab=256k, hidden=10k, output_layer in fp32: saves 10GB e2e
  • vocab=128k, hidden=10k, output_layer in fp32: saves 5GB e2e
image
MTP-depth-1 MTP-depth-2
iteration 1 pop → empty, alloc A pop → empty, alloc B ← the one new allocation
iteration 2+ pop A pop B — no new allocation

Problem: The step up to 2 never comes back down, because the pool never shrinks.

Fix: When a weight gets several backwards per iteration (MTP's repeated block), get_wgrad_tensor now drains that weight's own in-flight reduce-scatter and reuses its send buffer instead of allocating a second one, saving 10GB (vocab-256k) at no measurable throughput cost.

What

The GTP symmetric reduce-scatter send pool never shrinks, so every buffer it
holds is live at the memory peak. A weight with more than one backward per
iteration reaches get_wgrad_tensor while its previous send is still in flight,
so the pool grows by one permanent buffer — 5 GiB for a vocab-sized fp32
wgrad.

Found via MTP's repeated block, but not specific to it.

How

Drain that weight's own reduce-scatter and reuse the buffer instead of
allocating, trading one overlap for the permanent bytes.

  • RegisteredLIFOPool.has_free() limits the drain to the case where the
    alternative is a new allocation — no overlap given up otherwise.
  • rs_event.wait() fences the compute-stream wgrad GEMM against the
    still-reading send.
  • Skipped under CUDA-graph capture (host-side data-dependent branch).

Measured (rank 0, 4x GB200, MTP on)

before after
symmetric RS send pool 17.576 GiB 12.185 GiB
max allocated 133.659 GiB 128.269 GiB

−5.39 GiB, three independent measures agreeing; exactly two buckets lose a
duplicate. Loss bit-identical on iterations 1–2.

Scope

No-op without --gtp*-remat-nccl-ub, under CUDA-graph capture, or when a weight
has a single backward per iteration.

Tests

Two new has_free cases in test_gtp_symmetric_memory.py.

Convergence Test: passed on NT4 nano model

image

@fanshiqing
feat(gtp): Reuse the GTP wgrad send buffer across a weight's repeated…
11fb9c2 
… backwards

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
@fanshiqing
fanshiqing requested review from a team as code owners September 18, 2026 04:59
@copy-pr-bot

copy-pr-bot Bot commented Sep 18, 2026

Copy link
Copy Markdown

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.

@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft September 18, 2026 04:59
@github-actions

github-actions Bot commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

  1. Add the oncall reviewer (optional reviewer)
  2. Add required review teams based on your changes

See the contribution guide for more details.

@fanshiqing
fanshiqing marked this pull request as ready for review September 18, 2026 04:59
@fanshiqing fanshiqing self-assigned this Sep 18, 2026
@fanshiqing fanshiqing changed the title feat(gtp): Reuse the GTP wgrad send buffer across a weight's repeated backwards fix(memory): Reuse the GTP wgrad send buffer across a weight's repeated backwards Sep 18, 2026
@fanshiqing

fanshiqing commented Sep 18, 2026

Copy link
Copy Markdown
Member Author

/claude reivew

@fanshiqing

fanshiqing commented Sep 18, 2026

Copy link
Copy Markdown
Member Author

/ok to test 11fb9c2

@fanshiqing fanshiqing linked an issue Sep 18, 2026 that may be closed by this pull request
GTP memory optimizations #6609
Closed
@balasaajay balasaajay mentioned this pull request Sep 19, 2026
Closed
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Sep 19, 2026
@fanshiqing fanshiqing changed the title fix(memory): Reuse the GTP wgrad send buffer across a weight's repeated backwards fix(MTP's memory): Reuse the GTP wgrad send buffer across a weight's repeated backwards Sep 19, 2026
@fanshiqing fanshiqing linked an issue Sep 19, 2026 that may be closed by this pull request
[feat] GTP+MTP #6078
Closed
@fanshiqing fanshiqing mentioned this pull request Sep 19, 2026
Closed
@Phlip79

Phlip79 commented Sep 21, 2026

Copy link
Copy Markdown
Member

/review model=claude

@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 Sep 21, 2026
@fanshiqing
fanshiqing added this pull request to the merge queue Sep 21, 2026
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Sep 22, 2026
@fanshiqing @xuwchen
feat(gtp): Reuse the GTP wgrad send buffer across a weight's repeated…
c843425 
… backwards

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

(cherry picked from PR NVIDIA#7470 @ 11fb9c2; not merged to main at pick time)
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Sep 22, 2026
@fanshiqing @xuwchen
feat(gtp): Reuse the GTP wgrad send buffer across a weight's repeated…
b279511 
… backwards

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

(cherry picked from PR NVIDIA#7470 @ 11fb9c2; not merged to main at pick time)
@nemo-automation-bot

nemo-automation-bot Bot commented Sep 22, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit ed4320a Sep 22, 2026
114 of 117 checks passed
@fanshiqing
fanshiqing deleted the memory_opt/reuse_gpt_wgrad_send_buf_across_repeated_bwds branch September 22, 2026 09:42
@lhb8125 lhb8125 mentioned this pull request Oct 1, 2026
Open
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Oct 10, 2026
@fanshiqing @xuwchen
feat(gtp): Reuse the GTP wgrad send buffer across a weight's repeated…
474f71e 
… backwards (NVIDIA#7470)

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

(cherry picked from PR NVIDIA#7470 @ 11fb9c2; not merged to main at pick time)

Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Oct 14, 2026
@fanshiqing @xuwchen
feat(gtp): Reuse the GTP wgrad send buffer across a weight's repeated…
0951cdc 
… backwards (NVIDIA#7470)

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

(cherry picked from PR NVIDIA#7470 @ 11fb9c2; not merged to main at pick time)

Signed-off-by: xuwenc <xuwenc@nvidia.com>
fanshiqing added a commit to fanshiqing/Megatron-LM that referenced this pull request Oct 19, 2026
@fanshiqing
feat(gtp): Reuse the GTP wgrad send buffer across a weight's repeated…
21460b6 
… backwards (NVIDIA#7470)

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

(cherry picked from PR NVIDIA#7470 @ 11fb9c2; not merged to main at pick time)

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

This branch was previously deployed

2 inactive deployments
public — 11fb9c2b Deployed Sep 18, 2026 by copy-pr-bot[bot] via build-test-publish-wheels / publish-wheels (megatron-core, amd64) #9496
test — 11fb9c2b Deployed Sep 18, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #22036
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kvareddy kvareddy kvareddy approved these changes
@Phlip79 Phlip79 Phlip79 approved these changes
@deepakn94 deepakn94 Awaiting requested review from deepakn94

Assignees

@fanshiqing fanshiqing

Labels

Approved All necessary approvals have been made complexity: low nemotron

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

GTP memory optimizations [feat] GTP+MTP

4 participants

@fanshiqing @Phlip79 @kvareddy @svcnvidia-nemo-ci