[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 shortens CUDA graph tensor and closure lifetimes during warmup, capture, and teardown.

• Releases warmup and capture outputs after their final use.
• Snapshots per-callable graph state rather than retaining shared outer containers.
• Makes repeated reset calls safe and adds output-lifetime coverage.

Confidence Score: 4/5

The PR is not yet safe to merge because delayed-weight-gradient callers can crash by invoking backward_dw after resetting the graphed callable.

Reset clears bwd_dw_graph while the backward_dw closure retains a true need_bwd_dw flag, leaving a reachable None.replay() failure on the exposed post-reset lifecycle path.

Files Needing Attention: transformer_engine/pytorch/graph.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Reduces retained graph state and makes reset idempotent, but leaves backward_dw able to replay a graph reference cleared by reset.
tests/pytorch/test_cuda_graphs.py	Adds focused lifetime and repeated-reset coverage, but does not exercise backward_dw after reset.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  C[Captured callable] --> B[backward_dw closure]
  C --> R[reset closure]
  B -->|need_bwd_dw is true| G[bwd_dw_graph.replay]
  R --> X[Reset graphs and release static state]
  X --> N[bwd_dw_graph becomes None]
  N -->|later backward_dw call| E[AttributeError]

Loading

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

[greptile-apps]

Reset breaks backward_dw

When a callable captured for delayed weight-gradient computation invokes backward_dw() after reset(), the unchanged need_bwd_dw flag causes the closure to call replay() on the cleared bwd_dw_graph, raising AttributeError during teardown.

Suggested change

	def backward_dw():
	if need_bwd_dw_graph.get(graph_idx, False):
	bwd_dw_graphs[graph_idx].replay()
	if need_bwd_dw:
	bwd_dw_graph.replay()
	def backward_dw():
	if need_bwd_dw and bwd_dw_graph is not None:
	bwd_dw_graph.replay()