[JF-D]

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

What does this PR do?

Fixes local CUDA-graph training when CheckpointWithoutOutput recomputation uses parameters from nested modules that are not visible in the outer autograd graph.

Problem

CheckpointWithoutOutput runs its recompute function during backward. Parameters
used only by that function do not have autograd edges from the outer forward
outputs.

CUDA-graph capture previously discovered params_to_backprop only by traversing
the autograd graph from the captured forward outputs. Consequently, a parameter
used exclusively by checkpoint recomputation could be omitted from the captured
backward inputs.

This caused two related problems:

1. The nested recompute backward accumulated the parameter gradient into
   param.grad, but torch.autograd.grad returned None for the corresponding
   outer parameter input. The gradient was therefore not transferred to
   param.main_grad.

2. Because the parameter was not passed as a synthetic input to the CUDA-graph
   replay node, replay had no autograd edge to its AccumulateGrad node. The
   normal DDP hook registered on that node did not fire, leaving the DDP bucket
   with an incomplete ready count.

This can result in missing gradients, incorrect gradient accumulation, or DDP
waiting for a bucket that never becomes ready.

Example:

def forward(self, x):
    checkpoint = CheckpointWithoutOutput()


    hidden = checkpoint.checkpoint(self.checkpointed, x)
    output = self.projection(hidden)

    checkpoint.discard_output_and_register_recompute(output)
    return output

self.projection.weight is used by the outer forward autograd graph, so it can be
discovered by traversing backward from output. In contrast, self.checkpointed.weight
is closed over by the checkpoint recompute function. Its computation is deferred until
the nested recompute backward and therefore has no AccumulateGrad edge in the
outer graph. Traversing only the captured forward outputs consequently discovers
projection.weight but misses checkpointed.weight, even though backward eventually
produces a gradient for it.

Solution

Automatically discover hidden checkpoint dependencies

During CUDA-graph warmup, CheckpointWithoutOutput bypasses checkpointing and
executes its recompute function with autograd enabled. Parameters used by nested
recompute modules are therefore visible through their AccumulateGrad nodes.

This change:

• Traverses the warmup autograd graph to discover these parameters.
• Filters the result to parameters recursively owned by the captured base module.
• Merges the warmup-discovered parameters with those reachable from the actual
  captured forward outputs.
• Preserves base_module.parameters() ordering so capture and replay use a stable
  positional input layout.
• Runs at least one internal warmup for training graphs, even when
  cuda_graph_warmup_steps=0, because one discovery pass is required for correct
  checkpoint handling.

This discovery is automatic and does not require checkpoint call sites or
individual modules to declare additional dependencies.

Capture gradients produced by nested recomputation

For a parameter hidden behind CheckpointWithoutOutput,
torch.autograd.grad(..., allow_unused=True) preserves its requested output slot
but returns None for that slot. Its nested backward instead writes the real
gradient to param.grad.

Backward capture now transfers that contribution into param.main_grad.

The returned wgrad and param.grad are handled independently. This is
important when the same parameter is used both outside and inside checkpoint
recomputation:

• wgrad represents uses visible to the outer autograd graph.
• param.grad represents uses performed by the nested recompute backward.

Both are legitimate contributions and must be accumulated exactly once.

Preserve normal DDP bucket readiness during replay

All discovered parameters are passed as synthetic inputs to
_CudagraphReplayNode.apply. For ordinary DDP parameters, this creates replay
autograd edges to the existing AccumulateGrad nodes where DDP registered its
hooks.

For checkpoint parameters whose gradients were already transferred from
param.grad to param.main_grad, replay sets
grad_added_to_main_grad=True before the normal DDP hook runs. The hook then:

• Skips adding the same gradient to main_grad a second time.
• Clears param.grad.
• Still calls register_grad_ready(), allowing the DDP bucket to be dispatched.

GTP parameters retain their existing explicit hook path in
_CudagraphReplayNode.backward.

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

[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.

[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.

[JF-D]

@claude strict-review

[JF-D]

/ok to test 0616f63

[kvareddy]

@jiemingz can you please take a look at this MR?

[JF-D]

@claude strict-review

[JF-D]

/ok to test cb7e656

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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