[PyTorch] Reduce CUDA graph memory retention

[buptzyb]

Description

Reduce avoidable GPU-memory retention across CUDA graph construction and teardown.

• Release warmup forward outputs as soon as their scheduled backward consumes them, and release inference outputs immediately.
• Drop buffer-reuse capture locals after the per-callable containers take ownership, allowing weak-referenced graph-pool buffers to be reused by later captures.
• Keep make_graphed_attribute_functions while snapshotting only per-callable graph state, and clear replay closure state when reset() is called.

The changes preserve warmup/capture order and public APIs.

Testing

• TransformerEngine pre-commit formatting and Python 3.10 compatibility hooks on the modified files
• Source-built TransformerEngine on one H100
• python -m pytest -q tests/pytorch/test_cuda_graphs.py -k "warmup_releases_consumed_outputs or inference_warmup_does_not_retain_outputs or reused_capture_buffers_release_outputs_after_backward or reset_releases_only_the_selected_callable or capture_time_hooks or interleaved_pipeline_parallelism"
• Result: 10 passed

Signed-off-by: Robin Zhang robinz@nvidia.com

[greptile-apps]

Greptile Summary

The PR reduces CUDA graph memory retention by releasing warmup outputs promptly, dropping capture-local tensor references, and clearing each callable’s replay state during reset.

• Releases training warmup outputs immediately after their corresponding backward and inference outputs after forward.
• Transfers ownership of capture tensors to per-callable containers without retaining redundant local references.
• Makes reset idempotent and terminal across forward, backward, and delayed weight-gradient replay entry points.
• Adds focused CUDA lifecycle and weak-reference tests for output release, callable isolation, and post-reset behavior.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Releases temporary CUDA graph references earlier and implements isolated, idempotent, terminal reset semantics across all replay entry points.
tests/pytorch/test_cuda_graphs.py	Adds targeted GPU lifecycle tests covering warmup and capture output release, reset isolation, repeated reset, and rejected post-reset replay.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  F[Captured callable forward] --> A[Autograd replay state]
  A --> R[reset called]
  R --> G[Reset CUDA graphs]
  G --> C[Clear per-callable tensors and graph references]
  C --> T[Mark callable terminal]
  T --> E[Later replay entry point]
  E --> X[Raise controlled RuntimeError]

Loading

_{Reviews (3): Last reviewed commit: "[PyTorch] Release per-callable state on ..." | Re-trigger Greptile}

[janbernloehr]

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:

	stock graph.py	with PR #3427
deferred capture	OOM in _make_graphed_callables warm-up	completed in 20.9 s (rank 0)
training iterations	died at capture after iteration 3	30 / 30 completed
exit	crash (SIGABRT after OOM on multiple ranks)	clean exit 0
throughput	n/a	264,944 tokens/s aggregate (4,139.8 tokens/s/GPU)

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.