Skip to content

fix: ensure GTP params are ready before CUDA graph replay - #6525

Merged
fanshiqing merged 1 commit into
NVIDIA:mainfrom
JF-D:jiangfeid/gtp_cg_replay_param_readiness
Aug 21, 2026
Merged

fix: ensure GTP params are ready before CUDA graph replay#6525
fanshiqing merged 1 commit into
NVIDIA:mainfrom
JF-D:jiangfeid/gtp_cg_replay_param_readiness

Conversation

@JF-D

@JF-D JF-D commented Aug 13, 2026 •
edited by fanshiqing
Loading

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

What does this PR do?

PR #6388 ensures that GTP eager prefetch publishes DDP-sharded parameters
before reading them. During local CUDA graph capture, that publication must
not be captured because DDP parameter all-gathers are scheduled once per
iteration, while a graph may replay once per microbatch.

This PR carries the parameter-readiness dependency across the CUDA graph
capture/replay boundary:

  • Records parameters accessed by GTP all-gather during forward graph capture.
  • Deduplicates the recorded parameters for each graph.
  • Stores the dependency list on _CudaGraphRunner.
  • Calls ensure_params_ready() immediately before graph replay.
  • Clears the retained references when CUDA graphs are deleted.
  • Tracks GTP communication during forward graph capture, allowing the
    registrations to reach the graph runner.

Issue tracking

Related to #6388 and #6232.

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

@JF-D
JF-D requested review from a team as code owners August 13, 2026 22:06
@JF-D JF-D self-assigned this Aug 13, 2026
@JF-D JF-D added the nemotron label Aug 13, 2026
@JF-D

JF-D commented Aug 13, 2026

Copy link
Copy Markdown
Contributor Author

/claude strict-review

@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft August 13, 2026 22:09
@github-actions

github-actions Bot commented Aug 13, 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.

@copy-pr-bot

copy-pr-bot Bot commented Aug 13, 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 August 13, 2026 22:16
@JF-D
JF-D requested a review from fanshiqing August 14, 2026 05:25
@fanshiqing

fanshiqing commented Aug 14, 2026

Copy link
Copy Markdown
Member

/claude strict-review

@fanshiqing
fanshiqing requested a review from jiemingz August 14, 2026 10:47
@JF-D
JF-D force-pushed the jiangfeid/gtp_cg_replay_param_readiness branch from 3261774 to ad00e99 Compare August 14, 2026 21:04
@copy-pr-bot

copy-pr-bot Bot commented Aug 14, 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.

@JF-D
JF-D force-pushed the jiangfeid/gtp_cg_replay_param_readiness branch from ad00e99 to a74c65f Compare August 14, 2026 21:26
@JF-D
fix: ensure GTP params are ready before CUDA graph replay
876f098 
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
@JF-D
JF-D force-pushed the jiangfeid/gtp_cg_replay_param_readiness branch from 88df797 to 876f098 Compare August 20, 2026 01:29
@JF-D
JF-D enabled auto-merge August 20, 2026 01:30
@JF-D
JF-D added this pull request to the merge queue Aug 20, 2026
@kvareddy

kvareddy commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

@JF-D, let's wait until we have sign-off from Shiqing on GTP-related changes

@fanshiqing
fanshiqing removed this pull request from the merge queue due to a manual request Aug 20, 2026
@JF-D
JF-D added this pull request to the merge queue Aug 20, 2026
@nemo-automation-bot

nemo-automation-bot Bot commented Aug 20, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

@JF-D
JF-D removed this pull request from the merge queue due to a manual request Aug 20, 2026
@fanshiqing
fanshiqing added this pull request to the merge queue Aug 20, 2026
@nemo-automation-bot

nemo-automation-bot Bot commented Aug 21, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit b3c9100 Aug 21, 2026
91 of 92 checks passed
dimapihtar pushed a commit to dimapihtar/Megatron-LM that referenced this pull request Aug 21, 2026
@JF-D @dimapihtar
fix: ensure GTP params are ready before CUDA graph replay (NVIDIA#6525)
0040e76 
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
devnkong pushed a commit to devnkong/Megatron-LM that referenced this pull request Aug 22, 2026
@JF-D Kezhi Kong
fix: ensure GTP params are ready before CUDA graph replay (NVIDIA#6525)
03dce5e 
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
Signed-off-by: Kezhi Kong <kezhik@kezhik-mlt.client.nvidia.com>
@lhb8125 lhb8125 mentioned this pull request Aug 27, 2026
Open
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Aug 27, 2026
@JF-D @xuwchen
fix: ensure GTP params are ready before CUDA graph replay (NVIDIA#6525)
ba029eb 
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit b3c9100)
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Aug 27, 2026
@JF-D @xuwchen
fix: ensure GTP params are ready before CUDA graph replay (NVIDIA#6525)
f9b8388 
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit b3c9100)
ksivaman pushed a commit to ksivaman/Megatron-LM that referenced this pull request Sep 1, 2026
@JF-D
fix: ensure GTP params are ready before CUDA graph replay (NVIDIA#6525)
72dc163 
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
xuwchen pushed a commit to xuwchen/Megatron-LM that referenced this pull request Sep 3, 2026
@JF-D @xuwchen
fix: ensure GTP params are ready before CUDA graph replay (NVIDIA#6525)
d90f00e 
Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit b3c9100)

This branch was previously deployed

2 inactive deployments
public — 876f098b Deployed Aug 20, 2026 by copy-pr-bot[bot] via finalize / publish-docs #6716
test — 876f098b Deployed Aug 20, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #19551
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@fanshiqing fanshiqing fanshiqing approved these changes
@ericharper ericharper ericharper approved these changes
@jiemingz jiemingz jiemingz approved these changes

Assignees

@JF-D JF-D

Labels

Approved All necessary approvals have been made complexity: low

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[feat] GTP+CG

6 participants

@JF-D @fanshiqing @kvareddy @ericharper @jiemingz @svcnvidia-nemo-ci