[JF-D]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Problem

MTP can use the same GTP parameter multiple times in one forward/backward pass. Each completed GTP wgrad reduction contributes to main_grad and produces one DDP grad-ready
occurrence. DDP learns this per-parameter occurrence count during eager execution and waits for the same count in later iterations before dispatching the bucket reduction.

Local CUDA graph replay does not execute the Python GTP finalization path directly. Previously, it inferred finalized parameters from the captured graph structure and
deduplicated them by parameter identity. If a parameter was finalized twice, replay reported it ready only once, leaving DDP waiting for the second occurrence and
preventing the bucket reduction from being dispatched.

Solution

Record every actual GTP wgrad-finalization occurrence during backward graph capture, at the point where GTP invokes the DDP grad-ready hook.

The captured occurrence list intentionally preserves duplicate parameters. After graph replay, the hook plan reproduces the same number of grad-ready notifications as
eager execution, allowing DDP’s per-parameter ready counts to reach their expected values.

This also removes the previous graph-structure-based inference, so replay bookkeeping reflects the finalizations that actually occurred during capture.

Validation

• Added a unit test verifying that capture preserves repeated finalization occurrences such as [param_a, param_b, param_a].
• Full 4-GPU GTP unit-test suite: 159 passed, 8 skipped.

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

[github-actions]

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]

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]

/claude strict-review

[fanshiqing]

/claude review

[claude]

Light review — the core change (record actual finalization occurrences at the hook site instead of inferring them from prev_w/next_w graph structure) looks correct, and it is strictly more faithful than the old inference: whichever code path fires _grad_accum_hook during capture is the same path that fires it in eager, so replay counts match eager counts by construction. Two things worth a look.

1. finalized_params is no longer filtered by chain.

Everywhere else in create_bwd_graph selects GTP work by chain — capture_comms.get_comms_for_chain(GTPChain.GRAPHED.value) — precisely because, per track_gtp_capture_comms' own docstring, the tracker records ownership, not chain membership, and "nested modules can issue work for UNGRAPHED parameters while the context is active." The new line takes the list unfiltered:

self.finalized_during_bwd_capture = (
    list(capture_comms.finalized_params) if self.gtp_remat else []
)

The old _compute_finalized_during_bwd_capture walked self.params_to_backprop, so it was effectively same-chain. Now, if a captured module's backward finalizes an UNGRAPHED param (e.g. .mlp.shared_experts. under moe_shared_expert_overlap, which _classify_param_chain sends to UNGRAPHED while the enclosing layer is captured), that param lands in the hook plan and its hook fires on get_rs_stream(GTPChain.GRAPHED.value, group) — not its own chain's RS stream. Could you confirm that is intended? If not, the fix is a chain filter:

self.finalized_during_bwd_capture = (
    [
        p
        for p in capture_comms.finalized_params
        if getattr(p, "chain_id", None) == GTPChain.GRAPHED.value
    ]
    if self.gtp_remat
    else []
)

2. The new test does not cover the bug that was fixed.

test_wgrad_finalization_preserves_repeated_occurrences only exercises GTPCaptureCommState.finalized_params as a list — it asserts that three append calls produce three entries. That passes even if _handle_megatron_grad_accum never calls register_capture_wgrad_finalize, which is the actual fix. Two cheap additions that would pin the real behavior:

• That _handle_megatron_grad_accum registers one occurrence per call when _grad_accum_hook is set (and none when it is not). Doable with a stub param plus a monkeypatched get_dummy_wgrad, in the same non-distributed style as the sibling tests in this class.
• That a repeated param survives into _gtp_finalize_hook_plan with multiplicity — the dedup that caused the hang was downstream of the list, and the params_by_group grouping is where a future dedup would silently reintroduce it.

tests/unit_tests/generalized_tensor_parallel/test_gtp_mtp.py::TestGTPMTP::test_mtp_shared_weight_ddp_grad_ready_counts is the eager analogue of exactly this invariant; a CUDA-graph counterpart of it would be the strongest regression guard, though I understand that is a heavier lift.

[JF-D]

/ok to test e3b9bfe

[JF-D]

/claude strict-review

[copy-pr-bot]

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.

[nemo-automation-bot]

🔄 Merge queue validation started!

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