[htesd]

What does this PR do?

CheckpointFunction.backward in megatron/core/tensor_parallel/random.py 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

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.

[htesd]

For the record: this PR was not closed by me on purpose. My fork left the upstream fork network on Aug 5 (repo turned private for unrelated internal use), and GitHub auto-closes open PRs from a detached fork, attributing the close to the repo owner — hence the misleading "htesd closed this". It took #6017 down in the same second.

Re-opened as #6339 from a clean fork, on current dev (c76ff61), byte-identical change. Sorry for the noise.