[xiaoyao0115]

Description

Megatron's dynamic context parallelism (NVIDIA/Megatron-LM#6479) captures one CUDA graph bank per context-parallel (CP) size, because CP communicators are capture-time constants. Giving each bank its own memory pool multiplies graph memory by the number of banks.

This PR adds a private _graph_memory_slots argument to make_graphed_callables. It lets mutually exclusive graph variants share one torch.cuda.MemPool. FP8 attention and fine-grained activation offloading also work inside these graphs:

• FP8 attention scale snapshots and THD all-gather metadata now reach backward through saved tensors, which the slot arenas protect.
• A private _graph_memory_offload_context lets a captured scope offload selected saved tensors to pinned host arenas.

Without _graph_memory_slots, graph capture behaves as before, apart from the error-path cleanup (change 3) and the make_weak_ref fix (change 5). Two other changes also apply outside graphs, including eager training:

• FP8 attention and THD all-gather CP pass their backward state through prepare_for_saving (change 7). Under current scaling the CP paths now also snapshot the S scale, as FusedAttnFunc already did; other values are unchanged.
• Quantized wrappers are saved as physical buffers while v1 CPU offloading is active (change 2).

The branch is rebased onto current main (1f251a1) and has six commits, listed below.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue): make_weak_ref strides (change 5); FP8 wrappers saved under v1 CPU offload hooks (change 2)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

1. Slot memory (graph.py)

Each graph input carries (saved_arena, io_slot, branch, chunk, layer, warmup_group, grad_arena):

• saved_arena holds the tensors the forward saves for backward, plus its public outputs. grad_arena holds the gradients returned for user inputs. The caller (Megatron) assigns arenas so that two inputs share one only if their forward-to-backward lifetimes never overlap in any schedule it replays.

• Inputs with equal (io_slot, chunk, layer) are branches of one slot. They share output and gradient storage, so each arena is sized by the largest branch rather than the sum. Branches must:

  • name the same saved_arena and grad_arena;
  • expose identical output and gradient layouts (shape, stride, dtype, device, requires_grad, addressed bytes and alias topology). Outputs that view module parameters or buffers stay in the module's storage, but their layouts must match too.

  make_graphed_callables rejects branches that break either rule.

• Warmup runs once per warmup_group.

Forward replays snapshot saved tensors into the arenas, so a replay of another variant cannot overwrite them before the matching backward. Native saved tensors are grouped by storage identity, and aliases and versions are preserved. Public outputs that may overlap internally, such as expand views, are copied by their addressed storage span. The gradient passes through to the source view unchanged.

Contract: only tensors saved for backward (ctx.save_for_backward or saved-tensor hooks) are protected. For that reason, delayed weight-gradient _order entries are rejected when slots are used. Fused-wgrad parameters remain autograd inputs on replay. When a parameter lacks grad_added_to_main_grad, its placeholder gradient is zeroed, because autograd accumulates it into .grad.

2. Quantized tensors (quantized_tensor.py)

prepare_for_saving decomposes quantized wrappers (QuantizedTensorStorage, GroupedTensorStorage subclasses) into their physical storages in two cases:

• during slot capture;
• while v1 CPU offloading is active (NVTE_CPU_OFFLOAD_V1 and is_cpu_offload_enabled()).

restore_from_saved rebuilds the wrappers in backward. A hook that copies a wrapper by its nominal dtype would otherwise lose the FP8 data. All other paths save as before.

3. Capture error paths (graph.py)

If capture fails, GC state, temporary hooks, allocator settings and .grad state are restored. This also applies without slots.

4. Debug guard

With NVTE_CUDA_GRAPH_SLOT_DEBUG=1, each replay checks on the host that no other live graph owns the arena or output family it is about to reuse. The variable is documented in docs/envvars.rst.

5. make_weak_ref (utils.py)

make_weak_ref now keeps the strides of non-contiguous views. Previously, transposed or expanded static outputs were reinterpreted as row-major after the reuse block dropped their strong references.

6. Activation offloading in slot graphs (graph.py)

_graph_memory_offload_context(checker, *, d2h_stream=None, h2d_stream=None) is a private context manager. Inside a captured scope, the existing slot saved-tensor hook offloads the saves that checker selects, so no second hook pair is installed.

• Selected saves go to pinned host arenas. These are keyed by the same saved-arena IDs and sized by the largest variant. Offloaded saves do not also take space in the GPU saved arenas.
• The forward graph captures the D2H copies on d2h_stream and joins them before the forward ends. The backward graph reloads each storage group on h2d_stream when it is first unpacked.
• If any saved view of a storage stays on GPU, every view of that storage stays on GPU. Reloaded views keep their aliases.
• Warmup records which saves are offloaded, and capture follows that record. The caller owns the policy.
• When d2h_stream differs from the current stream, each offloaded save is first copied (its addressed span) on the current stream, and the D2H reads that copy. Without the copy, a later in-place write in the same forward could change the values before the D2H reads them.

7. FP8 attention state in saved tensors (backends.py, context_parallel.py)

Backward used to read FP8 scale snapshots and THD all-gather metadata from plain ctx attributes. Slot graphs do not protect those, so another variant's replay could overwrite them.

• FusedAttnFunc saves the S quantizer scale (Float8Quantizer) through prepare_for_saving and rebinds it in backward.
• CP P2P, all-gather and A2A save their scale snapshots through _prepare_backward_quantizers, and _restore_backward_quantizers rebinds them in backward. Like FusedAttnFunc, they snapshot every Float8Quantizer scale: QKV/O/S under delayed scaling and S under current scaling, whose S/dP use a delayed sub-recipe. Before, the CP paths snapshotted only under a delayed global recipe, so CP backward under current scaling read the live S scale.
• THD all-gather saves the per-step Q cu_seqlens, their padded variant, and cu_seqlens_kv_padded. The saved layout always holds two steps. NO_LOAD_BALANCE computes only one step, so it saves None for the second.

Possible split

If reviewers prefer, changes 3 (error-path cleanup), 5 (make_weak_ref), 2 (quantized decomposition), 7 (attention state) and 4 (debug guard) can each become a separate PR ahead of changes 1 and 6.

Commits

On main 1f251a1:

1. 83d11ca3 [PyTorch] Share CUDA graph memory across graph variants with slot arenas: changes 1–5.
2. 485b1a8e Protect slot graph outputs and captured attention state: changes 6 and 7, the v1-offload part of change 2, overlapping public outputs and full branch-layout checks (change 1), and tests/pytorch/attention/test_saved_quantizer_scales.py (wired into L0 and L1).
3. 9fc2372f Snapshot saved activations before asynchronous graph offload (change 6).
4. 58c37155 Keep the two-step THD Q metadata layout for NO_LOAD_BALANCE all-gather CP. This fixes a NO_LOAD_BALANCE THD all-gather backward failure that commit 2 introduced, with or without graphs.
5. ec656d35 Require slot branches to share arenas and document async offload copies: the arena-ID check (change 1), docstrings and pylint fixes.
6. 15fe37e0 Check layouts of module-state outputs across slot branches (change 1).

Review follow-up

• Expanded public outputs failed capture (r4370755121): 485b1a8e.
• Branch layouts were compared without shape or stride (r4370755124): 485b1a8e, extended to module-state outputs in 15fe37e0.
• Asynchronous D2H raced with later in-place writes (r4535182199): 9fc2372f.
• Branches of one slot could name different arenas (r4535799934): ec656d35.
• CI pylint C0116/W0613 on _CopyOverlappingOutput: ec656d35.

Testing

Tests in this PR:

• tests/pytorch/test_cuda_graphs.py adds slot tests. They cover:

  • variants sharing one backing pool;
  • liveness across PP/VPP schedules and slot wrap;
  • aliasing and staging contracts;
  • quantized (MXFP8, grouped) saved tensors;
  • returned and fused-wgrad parameter gradients;
  • capture-error cleanup;
  • strided and expanded outputs;
  • rejection of delayed wgrad, of mismatched branch layouts (including module-state outputs) and of mismatched branch arenas;
  • offloading: sync and async, mixed branches with slot wrap, and saved-view aliases with a mixed GPU/host policy;
  • mutated saved storage with offload none/sync/async. The async case delays the D2H stream with torch.cuda._sleep.

• tests/pytorch/attention/test_saved_quantizer_scales.py (new) covers:

  • CP scale snapshots and the fused-attention S scale restored from saved tensors;
  • FP8 wrapper buffers round-tripping through a v1 offload hook;
  • test_fp8_attention_slot_replay_matches_eager_with_overlapping_frames, which runs FP8 DPA with CP1 and CP2 banks sharing slots. Frames overlap within each bank (slot wrap) and across banks, and it covers {delayed, current scaling} × {SBHD, THD} × {no CP, P2P, A2A, all-gather} × offload off/on against eager.

  The CP cases need 2 or 4 ranks; test_cp_cases_with_two_ranks launches them under torchrun. L0 runs the single-process cases and L1 runs the CP cases. THD cases need SM100+ and cuDNN 9.23+.

• The NO_LOAD_BALANCE layout (commit 4) is covered by the existing tests/pytorch/attention/test_attention_with_cp.py::test_cp_with_{fused,flash}_attention_no_load_balance (L1).

Results at 15fe37e0 (NGC PyTorch 26.03 containers, cuDNN Frontend 1.32.2):

Selection	H100 (sm_90)	GB200 (sm_100)
test_cuda_graphs.py, whole file with the L0 environment (run twice)	540 passed, 426 skipped	781 passed, 1121 skipped
test_saved_quantizer_scales.py, L0 command (-k "not two_ranks")	9 passed, 30 skipped	15 passed, 24 skipped
test_saved_quantizer_scales.py, L1 command (two-rank CP cases)	4 passed, 20 skipped	24 passed
test_saved_quantizer_scales.py, 4 ranks	4 passed, 20 skipped	24 passed
test_attention_with_cp.py -k "all_gather or no_load_balance"	7 passed	79 passed
-k preserves_versions_of_mutated_saved_storage with graph.py before 9fc2372f	only [async] fails; all pass at head	only [async] fails; all pass at head
-k fused_attention_no_load_balance with context_parallel.py before 58c37155	fails; passes at head	fails; passes at head
test_cpu_offloading.py, test_cpu_offloading_v1.py, test_quantized_tensor.py, test_distributed_weight.py, test_attention.py -k fp8, test_kv_cache.py, test_attention_with_cp.py -k fp8, test_torch_compile.py -k dpa_torch_compile, compared test by test with main 1f251a12	no failure that main does not also have	same

On H100, the FP8 THD cases skip for hardware reasons. Black 24.4.2 (repository pre-commit arguments) and pylint 3.3.1 (pylintrc) pass locally on the changed files. CI has not run on this head yet.

test_slot_memory_variants_reuse_native_saved_allocations compares saved-arena sizes. An earlier version read the pool through torch.cuda.MemPool.snapshot(include_traces=False), which corrupted the Python heap in this container and made the whole-file run segfault later; the library code never calls it.

Megatron validation on this TE build: Megatron-LM#6479 9da6607dd (table in that PR). End-to-end throughput and memory have not been measured on this head; the earlier numbers predate the rebase and the FP8 and offload changes.

Limitations

• FP8 attention in slot graphs is tested with delayed and current scaling. MXFP8 attention has no S quantizer to snapshot and is not covered here.
• The offload policy is recorded at warmup, and a storage group stays on GPU if any of its views does. Megatron-LM#6479 freezes the policy per module scope, so graphs may offload less than eager.
• There is no ahead-of-use prefetch. Backward reloads a storage group when it is first unpacked and waits for it.
• With a distinct d2h_stream, each offloaded save costs one extra device-to-device copy on the compute stream. The copies stay alive until the end-of-forward join, so the forward graph's pool peak can grow by up to the sum of the offloaded spans. This has not been benchmarked.
• Input staging keeps its slot lifetime; offloading does not shorten it.
• Megatron-LM#6479, not TE, rejects two combinations under dynamic-CP graphs:
  • whole-layer capture with fine-grained activation offloading (partial cuda_graph_modules scopes are required);
  • pipeline parallelism with --dataloader-inter-document-masking.
• FP8 THD needs SM100+ and cuDNN 9.23+, so it is covered only on Blackwell.
• Megatron-LM#6479 needs a TE that contains all six commits. Its capability check looks only for _graph_memory_slots, plus _graph_memory_offload_context when offloading. It therefore accepts a TE that has only the first commit, which lacks the FP8 attention and THD all-gather protection.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The changes affect PyTorch CUDA graph capture, especially callers using shared slots.

Fix gradient copying for weights shared between fused and ordinary operations before merging.

Findings

1.  Shared weight gradients get overwritten ▶

Summary

Adds private CUDA graph slots so mutually exclusive CP variants can share one memory pool.

• Mutually exclusive graph variants can reuse caller-assigned memory slots.
• Quantized tensors and their scale data can travel through graph-saved tensors.
• Failed graph captures restore the temporary state they changed.
• Slot replays can check that no live graph still owns reused memory.
• Weak references keep the strides of non-contiguous tensor views.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[CP graph variants] --> P[Shared CUDA memory pool]
  P --> F[Forward replay]
  F --> S[Saved tensors and outputs in slot arenas]
  S --> B[Matching backward replay]
  B --> G[Returned gradients]
  B --> R[Slot available for reuse]

Loading

_{Reviews (4) · Last reviewed commit: "Check layouts of module-state outputs ac..." · Reviewed by Greptile}

[greptile-apps]

Shared weight gradients get overwritten

With slots enabled, a weight used by both a fused-wgrad TE layer and an ordinary PyTorch operation gets a real returned gradient from the ordinary operation. For example, a module can add self.linear(inp) to torch.nn.functional.linear(inp, self.linear.weight).

This check skips cloning that real gradient because the weight belongs to fused_wgrad_params, even with clone_param_grads_on_return=True. Replay returns the captured buffer directly, so the next backward overwrites retained gradients and can corrupt .grad accumulation. Skip cloning only for actual dummy gradients, not every weight used by a fused-wgrad layer.