[PyTorch] Share CUDA graph memory across dynamic CP variants

[xiaoyao0115]

Summary

• accept a compact seven-field slot/liveness plan in make_graphed_callables
• capture mutually exclusive dynamic-CP branches in one private torch.cuda.MemPool
• restore allocator checkpoints so equal PP/VPP physical slots reuse the same allocator topology and addresses across CP sizes
• preserve graph-boundary I/O, saved-tensor, parameter-gradient, and input-staging lifetimes with explicit liveness arenas
• retain relative byte offsets and physical overlap for aliased public output views, while leaving module-state views outside the slot pool
• validate every checkpoint owner before StorageImpl mutation, and transactionally restore the original allocator boundary if owner detach or checkpoint restore fails
• restore class wrappers, RNG, FP8, allocator settings, capture flags, GC, input-grad state, and temporary hooks on preparation or capture failure

Motivation

Dynamic context parallelism captures several CP-size alternatives. Capturing every alternative with independent graph allocations multiplies CUDA graph pool memory. Same-slot CP alternatives are mutually exclusive at replay time, so they can share the canonical branch's physical allocations.

This change keeps one private allocator pool and captures CP alternatives as checkpoint branches. Same liveness slots reuse the same addresses while retaining compatibility, saved-tensor interval, allocator-liveness, and StorageImpl ownership checks.

Compatibility and private contract

The slot-memory interface is private. It accepts exactly:

(saved_arena_id, physical_slot, io_branch_id, model_chunk, layer, warmup_alias_group, user_grad_arena_id)

The paired MCore PR constructs the scheduler-reachable plan and lockstep branch order. Both PRs must roll out together. The ordinary CUDA graph path remains unchanged when no slot plan is supplied.

Only PyTorch's native CUDA caching allocator is supported. Other backends, including cudaMallocAsync, are rejected before capture.

Companion change

• [dev] Share CUDA graph slots across dynamic CP variants Megatron-LM#6479 at commit 3de21ef4a70297bfeea1f7bec1f61771f094b81b

Validation

Final TE head: b07d7df77085964df35f5e3a2dd5678f7bc75d68.

• exact-final-head clean-process full split (Slurm 6431883): all 35 tests added by this PR passed; the parent surface reported 332 passed, 105 skipped, 35 deselected
• exact-final-head focused slot suite (Slurm 6427059): 29 passed, 443 deselected
• parent-head slot/non-slot split (Slurm 6413872): 22 passed, 443 deselected and 338 passed, 105 skipped, 22 deselected
• predecessor-head slot/output-alias suite and standalone probes (Slurm 6420861): 26 passed, 443 deselected; exact and nonzero-offset aliases, parameter/buffer aliases, gradients, capture-state cleanup, and tensor-subclass fail-fast all passed
• exact-final-head zero-layer PP/VPP regression (Slurm 6427059): trailing, middle, and consecutive-middle chunks with no graphable layers passed as part of the complete slot suite (29 passed, 443 deselected); the pre-fix control failed at graph.py:1991 with an out-of-range slot lookup (Slurm 6425845)
• three allocator failure points (partial owner detach, pre-setter failure, post-mutation setter failure) all rolled back, restored standard deleters, and allowed a fresh capture/replay in the same process (Slurm 6413150): 3 passed
• exact final source passes the CI-equivalent local PyTorch Python pylint==3.3.1 command with no diagnostics and 10.00/10; repository-exact local hooks pass
• GitHub Core passes. All and PyTorch ended only with the hosted-runner-lost-communication annotation and no downloadable job log; the available fork token cannot rerun those upstream jobs. JAX fails while building unchanged common/activation/gelu.cu because fatbinary cannot open gelu.compute_75.ptx.
• final-head Python compilation and git diff --check passed; every authored commit has a DCO trailer
• paired MCore final-source focused suite (Slurm 6413332): THD graph core 56 passed, 11 skipped, dynamic static-input selection 3 passed, router parent-group 1 passed, and packing 10 passed
• exact distributed MCore source (Slurm 6413335): complete THD file 65 passed, 2 skipped on both ranks; four-rank reroute/DCP and mixed-subgroup parent reduction passed on every rank
• mixed parent-transport graph on exact MCore and a runtime-equivalent TE predecessor (Slurm 6414410): CP4+CP2+CP1+CP1 replayed 50 times against one parent communicator
• exact-MCore-head dynamic-slot fuzz (Slurm 6414336, 6414384): 200 schedules / 2220 microbatches plus 800 scheduler-bound cases, with checkpoint-failure rollback passing in both runs
• exact MCore plus TE parent-head 10.59B PP4/VPP4 functional proxy (Slurm 6414248): CP1+CP2 on every training call, 16 graph/eager steps, zero skipped/NaN, exact sample and schedule match, maximum loss delta 0.002740 (0.02457%)
• intermediate-head 48-layer Qwen3-30B PP2/VPP4 multi-CP1/2 CUDA graph smoke (Slurm 6419837): capture, replay, training, validation, and test completed 0:0 with zero skipped/NaN; predecessor-head job 6421024 below is the authoritative 50-step Qwen3-30B control
• predecessor-head 48-layer Qwen3-30B PP2/VPP4 control (Slurm 6421024): fixed-CP2 graph/eager and multi-CP1/2 graph/eager each completed 50 steps with zero skipped/NaN; graph and eager LM loss matched exactly at every logged step within each CP policy; multi versus fixed had the same maximum 0.000040 (0.0003238%) difference in graph and eager controls. The final commit only skips schedule events whose chunks contain zero graphable layers; every Qwen chunk in this control was nonempty.
• corrected predecessor-source lifecycle control (Slurm 6424219): fixed captured 96 callables and multi captured 192; resetting every graph changed no allocator or whole-device counter immediately, callable release returned the graph-pool reservation, explicit-root release matched the root ledger, and the multi-minus-fixed allocator-external residual remained exactly +114 MiB/rank
• predecessor-source capture-phase attribution (Slurm 6424678): fixed and multi had the same initial +630 MiB/rank communicator warmup and zero external growth in the capture-time repeat; slot-bank preparation added zero external bytes; graph-bank capture added 150/148 MiB for fixed and 264/262 MiB for multi, an exact +114 MiB/rank multi-minus-fixed delta
• exact-final-head three-way input/output alias probe (Slurm 6429674): eager, native non-slot graph, and slot graph were compared. Both graph paths preserved output-output overlap, relative byte offset, write propagation, values, and gradients; neither graph path retained storage identity with the runtime input. Native graph retained an incidental alias to the capture sample, while slot graph intentionally moved that output to its arena.
• exact-final-head overlapping-liveness staging probe (Slurm 6430636): two liveness colors starting from one caller-owned sample storage received distinct staging targets and produced correct values and gradients

The complete test file also exposed a separate same-process runtime limitation. After hundreds of
parent TE graph captures, exact-head job 6430234 segfaulted while a newly constructed user
torch.cuda.MemPool was unwound. Unmodified TE base job 6431216 reproduces the same trigger
before any slot code exists. A pure PyTorch control retaining up to 512 graphs over 256 pool
handles passes, so the additional TE/runtime lifecycle condition is not yet identified. This is
why the exact PR regression result above uses clean processes; it does not claim arbitrary
same-process repeated TE capture lifecycles are proven safe.

Existing production-scale evidence predates the final audit commit: Qwen3-235B, PP4/VPP12, DCP CP4/8 completed 50/50 steps with no skipped or NaN iterations; DCP+CUDA graph versus eager DCP had maximum logged loss difference 0.00398 (0.0351%). This validates the core design but is not a final-head rerun.

In the predecessor-head Qwen3-30B control, slotized multi-CP1/2 versus slotized fixed-CP2 used one allocator pool in both cases and had a byte-identical explicit-root ledger on every rank. After step 50 and empty_cache(), the whole-device delta was +738 .. +832 MiB/rank (+0.4624 .. +0.5151%). Only +7.737 .. +7.807 MiB/rank remained allocated; reserved memory was +624 .. +718 MiB/rank, and the derived allocator-external delta was exactly +114 MiB/rank. Allocator snapshots place the reserved delta in inactive allocator topology, primarily the graph-private pool. Lifecycle and phase probes establish that the external +114 MiB/rank is created during capture of the additional CP graph bank, survives graph reset plus callable/root release, and is neither explicit tensor-root memory nor additional communicator initialization at the measured boundary. The available counters do not distinguish CUDA graph executable state from graph-capture-time driver or library retained/cache state, so no narrower object label is claimed.

Known limitations and review notes

• the Qwen3-30B 50-step control and graph lifecycle/capture-phase attribution are complete on the immediate predecessor head; the exact-final-head Qwen3-235B 50-step rerun is still pending
• allocator checkpointing and StorageImpl owner transfer use private PyTorch APIs
• owner validation occurs before mutation; failures after detach begins are rolled back to the original checkpoint. If that rollback itself fails, the process cannot safely continue
• expandable_segments restoration infers its original state from environment because PyTorch exposes no stable getter; prior programmatic allocator-setting changes can diverge
• the first positional sample argument and public CUDA outputs must be plain strided CUDA tensors; tensor subclasses, non-strided layouts, negative strides, and internally overlapping views are outside the supported contract
• output-output aliases and module-parameter/buffer views are preserved. Like native non-slot TE graph replay, slot replay does not retain storage identity between a public output and the runtime input; slot replay also does not retain native graph's incidental alias to the caller-owned capture sample because it moves the output into the slot arena
• liveness coloring is conflict-safe but not a weighted global optimum
• every CP alternative still has a separate CUDA graph executable and driver/library/communication metadata; only the MemPool and allocator/liveness topology are shared
• capture temporarily mutates process-wide class __call__, RNG, and allocator state and is not safe for concurrent calls
• the weak logical-CP to parent-transport map has no helper ownership token or refcount; overlapping
  helper lifetimes that share logical CP groups are unsupported because one teardown can remove the
  other's mapping
• the seven-field interface has no capability handshake, so mismatched MCore/TE versions fail rather than degrade
• creating a new user MemPool after a long sequence of existing TE graph tests can hit a baseline same-process segfault in the tested container; it reproduces without this PR and is not exercised by normal initial DCP capture