Fix CUDA graph parameter grad lifetime

[buptzyb]

Summary

Fix CUDA graph replay so parameter gradients returned from Graphed.backward do not expose CUDA graph static buffers to downstream autograd users.

The fix clones returned parameter gradients before handing them back to autograd, while preserving the existing aliasing behavior for delayed-wgrad parameters marked with skip_backward_post_hook.

Root Cause

When CUDA graph replay returns parameter grad slots directly from static graph buffers, downstream autograd users can retain references to buffers that are overwritten by later graph replays. This can corrupt retained grads or break gradient accumulation semantics.

This is related to PyTorch issue pytorch/pytorch#181723.

Changes

• Detect parameter grad slots in the graphed autograd input surface.
• Clone returned non-delayed-wgrad parameter grads before returning from Graphed.backward.
• Allow reused graph input/output buffer mode to weak-ref current parameter grad static buffers after capture because returned grads are now cloned.
• Add CUDA graph tests for owned returned parameter grads, accumulation, delayed-wgrad alias preservation, and reused buffer interleaved pipeline replay.

[greptile-apps]

Greptile Summary

This PR fixes a CUDA graph parameter grad lifetime bug where Graphed.backward returned direct references to static graph buffers, causing later graph replays to overwrite tensors retained by hooks or gradient accumulators. The fix clones returned parameter gradients at replay time, with a captured-at-build-time policy snapshot to avoid live attribute re-reads that could desync clone decisions from weak-ref decisions.

• _returned_param_grad_clone_slots is computed once at graph capture time by inspecting skip_backward_post_hook on each module parameter; delayed-wgrad parameters are excluded from cloning and the clone_param_grads_on_return knob lets callers opt out.
• In the _reuse_graph_input_output_buffers interleaved-pipeline path, parameter grad slots whose clone policy is True are now also weak-refed after capture (safe because the clone done at replay re-materializes the value before the weak ref is accessed), matching the existing write-ahead memory management pattern.
• Six new targeted tests cover: returned grad ownership, accumulation correctness, delayed-wgrad alias preservation, opt-out cloning, policy snapshot immutability, and reused-buffer interleaved pipeline with both clone modes.

Confidence Score: 5/5

Safe to merge. The change is narrowly scoped to the grad-return path in Graphed.backward and does not alter forward graph replay or any non-parameter grad handling.

The fix correctly isolates the clone policy decision to capture time, consistently applies it in both the interleaved-pipeline and standard code paths, and the weak-ref/clone interaction in the reuse-buffer path is sound because CUDA graph memory remains valid for the lifetime of the graph. All six new tests exercise the exact invariants the fix is meant to provide, and the previous review concerns about snapshot timing have been addressed.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Adds _returned_param_grad_clone_slots snapshot at capture time; clones parameter grad slots in Graphed.backward; weak-refs those slots in the reuse-buffer path. Logic is consistent between the interleaved and non-interleaved branches, and the policy is correctly frozen at capture.
tests/pytorch/test_cuda_graphs.py	Adds six new unit tests covering every combination of the new behavior: grad ownership, accumulation, delayed-wgrad alias, opt-out mode, policy snapshot, and reused-buffer interleaved pipeline with and without cloning. Existing pipeline test strengthened with final-weight equality check.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Graphed as Graphed.backward
    participant CUDAGraph as CUDA Graph static buffer
    participant Autograd as PyTorch Autograd

    Note over Caller,Autograd: Capture time
    Graphed->>CUDAGraph: Record backward kernel, grad written to static_grad_inputs[param_slot]
    Graphed->>Graphed: _returned_param_grad_clone_slots() snapshots skip_backward_post_hook per param
    Note over Graphed: Policy frozen at capture. Reuse path weak-refs param grad slots.

    Note over Caller,Autograd: Replay BEFORE this PR
    Caller->>Graphed: backward(grads)
    Graphed->>CUDAGraph: replay() overwrites static buffer
    Graphed-->>Autograd: return static_grad_inputs[param_slot].detach() aliases live buffer
    Autograd-->>Caller: param.grad points to static buffer
    Note over Caller: Next replay overwrites param.grad

    Note over Caller,Autograd: Replay AFTER this PR
    Caller->>Graphed: backward(grads)
    Graphed->>CUDAGraph: replay() overwrites static buffer
    Graphed->>Graphed: grad_input.detach().clone() for each param clone slot
    Graphed-->>Autograd: return owned clone independent tensor
    Autograd-->>Caller: param.grad is owned tensor, safe across replays

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant Graphed as Graphed.backward
    participant CUDAGraph as CUDA Graph static buffer
    participant Autograd as PyTorch Autograd

    Note over Caller,Autograd: Capture time
    Graphed->>CUDAGraph: Record backward kernel, grad written to static_grad_inputs[param_slot]
    Graphed->>Graphed: _returned_param_grad_clone_slots() snapshots skip_backward_post_hook per param
    Note over Graphed: Policy frozen at capture. Reuse path weak-refs param grad slots.

    Note over Caller,Autograd: Replay BEFORE this PR
    Caller->>Graphed: backward(grads)
    Graphed->>CUDAGraph: replay() overwrites static buffer
    Graphed-->>Autograd: return static_grad_inputs[param_slot].detach() aliases live buffer
    Autograd-->>Caller: param.grad points to static buffer
    Note over Caller: Next replay overwrites param.grad

    Note over Caller,Autograd: Replay AFTER this PR
    Caller->>Graphed: backward(grads)
    Graphed->>CUDAGraph: replay() overwrites static buffer
    Graphed->>Graphed: grad_input.detach().clone() for each param clone slot
    Graphed-->>Autograd: return owned clone independent tensor
    Autograd-->>Caller: param.grad is owned tensor, safe across replays

Loading

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into fix/cudagraph-w..." | Re-trigger Greptile}

[ptrendx]

I have a problem with this PR. On one hand I agree that there is a bug here and we should ensure that the gradient buffers are not overwritten before being applied. On the other hand for the cases without gradient accumulation (or when the accumulation is done in a different way, like in Megatron) it is a performance loss. Could we make this behavior optional - have it on by default, but also provide an opt-out option with a clear warning that states where this would be applicable?

[buptzyb]

@ptrendx I added a new argument _clone_param_grads_on_return (default = True). Documented that disabling this may improve performance but no longer have standard PyTorch returned-gradient lifetime semantics. Does this work for you?

[buptzyb]

Hi @ptrendx can we merge this PR?

[pggPL]

I think we should expose it to make_graphed_callables(), not just _make_graphed_callables() and also document it in docstring. Also, I think we can make it fully public without _.

[pggPL]

Also, I think we should make sure that MCore will set it to False if they do their wgrad accumulation. Maybe it's a good idea to create PR there also?

[buptzyb]

Thanks! Addressed both suggestions in f286621:

• clone_param_grads_on_return is now a public argument on make_graphed_callables() and is documented in the public docstring.
• I opened Disable TE CUDA graph parameter gradient clones Megatron-LM#5893 to set clone_param_grads_on_return=False from MCore when the public option is available, while preserving compatibility with older TE versions.

Thanks for the suggestions!

[pggPL]

/te-ci pytorch

[pggPL]

/te-ci pytorch