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

[yaox12]

/claude strict-review

[claude]

[IMPORTANT Correctness] Please confirm this is safe when backward runs under CUDA graph capture.

What: Before this change, the detached leaf's .grad was a second, independent strong owner of the input-gradient buffer. After it, the only owner is the grads tuple, which is handed to the autograd engine, consumed by the upstream node, and then freed.

Why it matters: backward() here explicitly tolerates running during capture — line 603 gates the _is_checkpoint_valid() error on not is_graph_capturing(), and cuda_graphs.py:1063 documents the [interaction with recompute] case where the recompute forward is re-run inside the backward pass. If torch.autograd.backward(outputs, args) on line 632 executes while a graph is being captured, the .grad tensors that AccumulateGrad produces are allocated from the capture mempool (self.mempool). Dropping the last Python reference to such a block returns it to that mempool, where a subsequent capture in the same pool can hand the same address to a different tensor — the two captured graphs then alias on replay and produce silently wrong gradients. This is precisely the hazard the create_strong_ref / fwd_cudagraph_buffer / bwd_cudagraph_buffer machinery in cuda_graphs.py exists to prevent, and why gc.collect() there is deliberately restricted to the last layer.

I believe the standard Megatron cudagraph path is unaffected, because checkpoint() (line 658) bypasses CheckpointFunction entirely under is_graph_warmup() or is_graph_capturing(), so no CheckpointFunction node ends up inside the captured forward graph and create_bwd_graph never captures its backward. But line 603 only makes sense if some configuration does reach this code during capture, so the invariant is not obvious from the file alone.

Suggested fix: either confirm via a run of the cudagraph + recompute_granularity=full functional tests (label the PR "Run functional tests" per .claude/skills/mcore-testing/SKILL.md) that this path is never entered during capture, or guard the release so it is skipped while capturing:

from megatron.core.transformer.cuda_graphs import is_graph_capturing
# (already imported at the top of this method)

# Release the recomputation-leaf gradients ... (comment)
# Skipped during graph capture: the returned `grads` are the only remaining
# owner, and freeing a capture-mempool block lets a later capture alias it.
if not is_graph_capturing():
    for inp in detached_inputs:
        if isinstance(inp, torch.Tensor):
            inp.grad = None

A one-line note in the PR description recording which configuration you validated would also be useful, since the reported measurement is from a non-cudagraph run.

[claude]

[SUGGESTION Other] The comment says the stale .grad is retained "across iterations", but the mechanism it describes — saved-variable references from the recomputation subgraph keeping the detached leaves alive — does not by itself explain an iteration-length lifetime.

The recomputed subgraph built at line 623 is dropped as soon as torch.autograd.backward on line 632 finishes (it is not retained), and ctx is torn down when the engine releases the node. What actually extends the leaves' lifetime past that point is whatever external reference the measurement found, e.g. an outer autograd node or a Python-side container still holding the recomputed tensors. Without naming it, a future reader cannot tell whether the release is still necessary or has become dead code, and cannot reproduce the finding.

The PR description says weakref instrumentation identified the retaining path. Naming that path here — the specific holder, not just "saved-variable references" — would make the comment verifiable:

# Release the recomputation-leaf gradients. They have already been harvested
# into `grads` above, so the returned gradients are unaffected. The detached
# leaves can outlive this call when <name the concrete holder>, in which case
# their `.grad` keeps a full copy of the segment's input gradients alive until
# the holder is dropped. Nested checkpointing is unaffected: each call touches
# only its own `detached_inputs`.

Also worth noting inline: this loop is safe only because no caller reads detached_inputs[i].grad after backward returns. I checked and none does — CheckpointWithoutOutputFunction.backward (line 766) reads inp.grad from its own ctx.inputs, and ScheduleNode.get_grad (pipeline_parallel/utils.py:290) reads from its own detached self.inputs, both distinct tensor sets. Stating that invariant in the comment protects it against a future change.

[claude]

Strict Review Summary

Findings: CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 1

Scope

Single hunk, 9 added lines, 0 deleted, in megatron/core/tensor_parallel/random.py. CheckpointFunction.backward clears .grad on the detached recomputation leaves after harvesting them into the returned grads tuple.

What I verified

• Returned gradients are unchanged. grads is built on line 633, before the release loop. It holds strong references to the gradient tensors, so setting inp.grad = None afterwards rebinds the leaf attribute only — it does not mutate or free the tensors that line 645 returns. The PR's claim here is correct.
• No caller reads these leaves' .grad afterwards. I traced every .grad read in the checkpoint/recompute/pipeline paths. CheckpointWithoutOutputFunction.backward (line 766) reads from its own ctx.inputs, and ScheduleNode.get_grad (megatron/core/pipeline_parallel/utils.py:290) reads from its own detached self.inputs — both are distinct tensor sets from this function's detached_inputs, which are local to this call and never published.
• Nested checkpointing is unaffected, as stated: detached_inputs is a call-local product of detach_variable(inputs), so an inner call cannot reach an outer call's leaves.
• No new identifiers introduced. The loop variable inp shadows the generator variable on line 633 but does not leak (that one is generator-scoped). Nothing added is unused.
• No process-group usage added, so the Megatron Core parallel_state guidance in CLAUDE.md does not apply.
• No API, config, checkpoint-format, or default-value change — backward compatibility is not at risk. Adding a .grad = None after the gradients have been collected has no numerical effect and no measurable throughput cost (a short Python loop once per checkpointed segment per backward, no device sync).

Most impactful finding

The one substantive concern is the CUDA graph capture interaction. This change removes the second strong reference to the input-gradient buffers, leaving grads as the sole owner. If this backward ever executes during graph capture, those buffers come from the capture mempool, and dropping the last reference returns the block for reuse by a later capture in the same pool — the aliasing hazard that create_strong_ref / fwd_cudagraph_buffer / bwd_cudagraph_buffer in cuda_graphs.py exists to prevent.

My reading is that the standard path is safe, because checkpoint() at line 658 bypasses CheckpointFunction entirely under is_graph_warmup() or is_graph_capturing(). But line 603 deliberately suppresses the _is_checkpoint_valid() error when is_graph_capturing() is true, and cuda_graphs.py:1063 documents an [interaction with recompute] case where the forward is re-run during backward — so the file itself implies some configuration reaches this code during capture. Worth confirming rather than inferring.

Assessment

Risk: low. The change is minimal, local, numerically inert on the gradient values, and the memory-lifetime reasoning is sound for the non-cudagraph path. The measured 0.27 GB saving is small and honestly reported, with the 2.4 GB figure correctly labelled an extrapolation rather than a measurement.

Two things before merge:

1. Confirm the cudagraph + recompute_granularity=full path — add the "Run functional tests" label per .claude/skills/mcore-testing/SKILL.md, or add the is_graph_capturing() guard from the inline comment.
2. Sharpen the code comment to name the concrete retaining reference the weakref instrumentation found, so the release can be re-validated (or retired) later.

Neither blocks the approach — this is a reasonable, well-scoped memory fix.