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