[PyTorch] Share CUDA graph memory across dynamic CP variants

[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, which lets mutually exclusive graph variants share one torch.cuda.MemPool. Without the argument, graph capture behaves as before, apart from the error-path cleanup (change 3) and the make_weak_ref fix (change 5).

The PR now targets main (it previously targeted release_v2.14) and is a single commit.

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)
• 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 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 must expose identical output and gradient layouts, and they share output and gradient storage, so each arena is sized by the largest branch rather than the sum.
• 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.

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)

During slot capture, prepare_for_saving decomposes quantized wrappers into their physical storages, so the arena path can place them; backward then rebuilds the wrappers.

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.

Possible split

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

Testing

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 outputs;
• the rejection of delayed wgrad.

This commit (on main 53e241c) was built from source in the torch 2.11 NGC container and tested with tests/pytorch/test_cuda_graphs.py. The DPA and FP8-weight-caching graph tests run in separate processes (see the note below the table).

Selection	H100 (sm_90)	GB200 (sm_100)
-k slot_memory	59 passed, 6 skipped	62 passed
test_make_graphed_callables_with_dot_product_attention	3 passed	3 passed
test_make_graphed_callables_with_fp8_weight_caching	180 passed	290 passed
all other tests	260 passed	388 passed

With upstream main's graph.py and quantized_tensor.py, the "all other tests" selection on H100 gives 255 passed and 5 failed. The 5 failures are this PR's new capture-error-path tests.

In this container, a process that has run these graph tests segfaults at interpreter exit, after pytest has written its report. This happens with upstream main's Python files too: the DPA and FP8 graph tests leave the Python heap corrupted, and a later allocation crashes.

Megatron validation (unit tests on this TE build, and the Qwen3.5 benchmark) is in NVIDIA/Megatron-LM#6479.

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