Fix cross-iteration gradient retention in CheckpointFunction.backward

[htesd]

What does this PR do?

Re-submission of #6117 on current dev (c76ff61). The change is byte-identical to the original — megatron/core/tensor_parallel/random.py is unchanged since that PR was opened, so it applies cleanly.

Why the original PR disappeared: #6117 was not withdrawn by me. My fork left the upstream fork network on Aug 5 (the repository was turned private for unrelated internal use), and GitHub auto-closes any open PR coming from a detached fork, attributing the close to the repo owner — which is why the timeline reads "htesd closed this". It took #6017 down in the same second; that one is re-submitted as #6338. This PR comes from a fresh, properly attached fork (htesd/Megatron-LM-oss).

#6117 never got past copy-pr-bot, so it has no prior review or CI history. @yaox12 — sorry to ping you directly, you handled the sibling PR; would you mind vetting this one with /ok to test so CI can run?

The bug

CheckpointFunction.backward recomputes the forward pass using detached segment inputs, runs backward on the recomputed graph, and collects the input gradients from the detached leaves:

grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)

The detached leaves can remain reachable after backward returns when the recomputed graph retains references to them. In that case, their .grad fields also keep the input-gradient storage alive longer than required.

This change clears each detached tensor's .grad after collecting the returned gradients. The grads tuple retains the gradient tensors, so the gradients returned by CheckpointFunction.backward are unchanged. Nested checkpointing is unaffected: each call only touches its own detached_inputs.

Measured impact

Numbers below are from #6117 (measured on that PR's base commit); the affected code path is unchanged on current dev.

On a single H200, a 2-layer GPT configuration with num_residual_streams=4, bf16, micro-batch size 1, mock data, and recompute_granularity='full' produced the following allocator floor, in otherwise identical runs differing only by this change:

metric	before	after
Between-iteration floor, sequence length 8192	28.63 GB	28.36 GB

This is a reduction of 0.27 GB for the measured configuration.

Weakref instrumentation confirmed that the detached recomputation inputs and their .grad storage remained reachable during the between-iteration interval before this change, and that clearing those .grad fields manually released the same allocator blocks. With this change, the gradient references are cleared immediately after collection.

Assuming linear scaling with checkpointed segment count and sequence length, the measured result corresponds to an estimated reduction of approximately 2.4 GB for a 9-layer pipeline stage at sequence length 16384. This figure is an extrapolation, not a direct measurement.

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