Greptile Summary

The PR reduces CUDA graph memory retention while making per-callable reset terminal, isolated, and idempotent.

• Releases warmup outputs immediately after their final use.
• Drops capture-time local tensor references after per-callable containers assume ownership.
• Snapshots lifecycle state per callable and consistently rejects replay after reset.
• Adds CUDA lifecycle tests for output release, callable isolation, repeated reset, and replay rejection.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; both previously reported reset defects are addressed by guarded, per-callable teardown and replay checks.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  C[Capture callable] --> O[Per-callable graph and static state]
  O --> R[Replay forward/backward]
  O --> X[reset]
  X --> G[Reset captured CUDA graphs]
  X --> S[Release static closure state]
  X --> T[Mark callable terminal]
  T --> E[Later replay raises RuntimeError]

Reviews (5): Last reviewed commit: "[PyTorch] Bundle per-callable lifecycle ..." | Re-trigger Greptile

Validation report: this PR fixes an OOM during deferred CUDA-graph capture at 64-GPU scale

We hit a reproducible torch.OutOfMemoryError during TransformerEngine's deferred CUDA-graph capture on a 64-GPU Llama-3 70B NVFP4 pre-training run, and this PR resolves it. We ran a controlled A/B in which the only difference between the two arms was the contents of transformer_engine/pytorch/graph.py. Sharing the evidence in case it is useful for review.

Environment and workload

• Hardware: GB200 NVL, 16 nodes x 4 GPUs = 64 GPUs (~184.3 GiB usable HBM per GPU)
• Stack: Megatron-LM + Megatron-Bridge pre-training, PyTorch 2.x, Python 3.12, TransformerEngine installed prebuilt into the image
• Model: Llama-3 70B, NVFP4 precision, sequence length 8192, global batch size 256, micro-batch size 1
• Parallelism: TP=2, PP=4, VP=5 (interleaved virtual pipelining), CP=1, EP=1
• CUDA graphs: cuda_graph_impl='transformer_engine', cuda_graph_modules=[mlp, attn], 20 graphable layers per rank
• Capture is deferred: three normal training iterations run first, then Megatron calls make_graphed_callables() for all layers at once

Because VP=5, the run takes the interleaved _order code path in _make_graphed_callables — the same path this PR's test_ordered_warmup_releases_consumed_outputs covers.

Failure without this PR

The run trains normally through iteration 3, then dies as soon as capture begins:

INFO:megatron.core.transformer.cuda_graphs:Start CUDA Graphs capture...
[rank11]: expandable_segments: memory mapping failed with OOM on device 3
          while trying to map 20971520 bytes (free: 15269888, total: 197897486336)
OutOfMemoryError: CUDA out of memory. Tried to allocate 224.00 MiB.
GPU 3 has a total capacity of 184.31 GiB of which 14.56 MiB is free.
Of the allocated memory 170.32 GiB is allocated by PyTorch, with 228.08 MiB
allocated in private pools (e.g., CUDA Graphs), and 8.48 GiB is reserved by
PyTorch but unallocated.

Traceback (identical in shape on every failing rank):

  File "megatron/core/transformer/cuda_graphs.py", line 2576, in create_cudagraphs
    graphs = make_graphed_callables(
  File ".../transformer_engine/pytorch/graph.py", line 1612, in make_graphed_callables
    graphed_callables = _make_graphed_callables(
  File ".../transformer_engine/pytorch/graph.py", line 719, in _make_graphed_callables
    outputs = func(*args, **kwargs)
  File ".../transformer_engine/pytorch/graph.py", line 1585, in call_func
    outputs = old_call_funcs[block_cls](self, *args, **kwargs)
  File "megatron/core/transformer/transformer_layer.py", line 1181, in _te_cuda_graph_capture
    hidden_states = self._forward_mlp(hidden_states)

So the OOM occurs in the warm-up forward inside _make_graphed_callables, not in steady-state training — the model itself fits and trains fine for three iterations at ~170 GiB resident. Note the allocator breakdown: only ~228 MiB is in CUDA-graph private pools at the point of failure, i.e. the memory pressure is from retained ordinary allocations during warm-up, not from the graph pools themselves.

Result with this PR applied

Same image, same 64-GPU allocation, same recipe, same flags. We applied only this PR's change to transformer_engine/pytorch/graph.py on top of the graph.py already installed in the image (it applied cleanly with git apply --check), and left everything else byte-identical:

Why this PR appears to address it

On the interleaved _order warm-up path, the pre-PR code kept every warm-up forward's flattened outputs alive in per_fwd_outputs for the whole warm-up loop, reading them with per_fwd_outputs[per_callable_bwd_idx] and never removing the entry. With 20 graphable layers per rank and VP=5, that is a large number of live activation sets accumulating simultaneously, on top of an already ~170 GiB resident model. This PR's change to per_fwd_outputs.pop(per_callable_bwd_idx) plus the del outputs after each warm-up backward releases each set as soon as its backward consumes it, which matches the observed behavior exactly: the failure is in warm-up, and it disappears with this change.

The capture-phase del static_outputs, static_grad_inputs, grad_inputs is likely also relevant at this scale, since it lets weak-referenced graph-pool buffers be reused across the 20 sequential per-layer captures.

Scope and caveats

• We applied only the transformer_engine/pytorch/graph.py portion of this PR. The changes to tests/pytorch/test_cuda_graphs.py were not needed at runtime and were not applied.
• This is one configuration (single model/precision/parallelism/scale), run once per arm. It is strong evidence that this PR fixes this failure mode, but it is not a broad regression sweep.
• We did not bisect which upstream commit introduced the retention, so we cannot say whether this is a recent regression or a long-standing limit that this scale is the first to cross.
• To keep the comparison honest we verified the image contained exactly one installed copy of graph.py on the interpreter path and recorded its sha256 before and after patching. Both arms started from an identical pre-patch hash; only the patched arm's hash changed. No other file, dependency, image layer, or recipe parameter differed between the arms.

Validated against this PR at commit a3f1d527a9fd879ba7133505814a289b3b2da9c7. Happy to re-run this workload against any later revision of the PR if that would help.

This issue was drafted with assistance from the opus AI model.

/te-ci pytorch