Skip to content

Fix cross-iteration gradient retention in CheckpointFunction.backward #6339

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
htesd wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

Fix cross-iteration gradient retention in CheckpointFunction.backward #6339

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions megatron/core/tensor_parallel/random.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -632,6 +632,15 @@ def backward(ctx, *args):
torch.autograd.backward(outputs, args)
grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)

# Release the recomputation-leaf gradients: they were harvested into
# `grads` above, but the detached leaves can outlive this call (e.g.
# kept alive by saved-variable references from the recomputation
# subgraph), retaining a stale copy of the segment's input gradients
# across iterations.
Comment on lines +635 to +639

@claude claude Bot Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

for inp in detached_inputs:
if isinstance(inp, torch.Tensor):
inp.grad = None
Comment on lines +640 to +642

@claude claude Bot Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


_unset_checkpointing()
return (None, None) + grads

Expand Down