[PyTorch] Share CUDA graph memory across dynamic CP variants

[xiaoyao0115]

Summary

• add a nine-field graph-memory slot plan to make_graphed_callables
• capture mutually exclusive dynamic-CP branches in one torch.cuda.MemPool
• restore allocator checkpoints so the same PP/VPP slot reuses the same addresses across CP sizes
• preserve graph-boundary I/O, saved-tensor, parameter-gradient, and input-staging lifetimes
• add focused coverage for native saved tensors, overlapping views, fused wgrad, and PP/VPP schedules

Motivation

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

This change keeps one allocator pool and reuses addresses at the same liveness slot while retaining compatibility, allocator-liveness, saved-tensor interval, and StorageImpl ownership checks.

Compatibility

The slot-memory interface is private and accepts only the production nine-field plan. The ordinary CUDA graph path is unchanged when no slot plan is supplied.

Companion change

• [dev] Share CUDA graph slots across dynamic CP variants Megatron-LM#6479 constructs the PP/VPP liveness plan and lockstep CP capture order.

Validation

Static checks on commit c095214af883862527b5e6d35d978acd9b272844:

• Python compilation and git diff --check
• full-repository pylint: 10.00/10
• pre-commit on all changed files: merge-conflict, large-file, EOF, trailing-whitespace, Black, clang-format, and vermin checks passed locally
• DCO and GitHub Build/Check/PyTorch C++/PyTorch Python/JAX C++/JAX Python checks passed

Focused GPU validation:

• exact-head PyTorch extension build passed (Slurm 6073801)
• slot-memory suite: 15 passed, 438 deselected (Slurm 6073868)

Full regression used 8 GB200 GPUs across 2 nodes with a Qwen3-30B-A3B-shaped 48-layer model: PP4, VPP4, EP2, sequence length 8192, global batch 64, 32 runtime microbatches, and pad-packed-seq-alignment=max-seqlen-per-dp-cp-rank=8192. The dynamic schedule exercised CP1 and CP2. Capture logs report 2 dynamic-CP variants over 8 slot-aliased physical TE pool slots, confirming one physical slot ring rather than one bank per CP size.

• DCP+CUDA graph versus eager DCP, 16 training steps (Slurm 6074897): no skipped or NaN iterations; maximum logged absolute LM-loss difference was 0.00398 (0.0351% of the corresponding loss), and auxiliary-loss values matched at logged precision.
• DCP+CUDA graph versus fixed-CP2+CUDA graph, two independent 3-step paired runs (Slurm 6074298, 6074899): after iteration 3 and torch.cuda.empty_cache(), per-rank whole-device deltas were -440 to +16 MiB and -366 to +96 MiB. The largest observed positive delta was 0.1001% of device usage. There is no multiplicative memory growth with two CP variants.
• The allocator view is not byte-identical to fixed CP: active allocation was repeatable at +80 to +272 MiB, while reserved-memory deltas ranged from +758 to +1296 MiB across both runs. This is reported separately from whole-device usage and remains visible for follow-up profiling.

CI notes

• The GitHub JAX aggregate job currently fails while compiling unmodified common CUDA code with fatbinary fatal: Could not open input file 'gelu.compute_75.ptx'; JAX C++ and JAX Python checks pass.
• pre-commit.ci currently runs the repository's vermin hook under Python 3.14 and crashes inside vermin on the removed ast.Str API. The same changed-file hook set passes locally.