perf(cuda-graphs): share partial capture pools

[yaoyu-33]

Summary

Stacked on #2932 and linked to #1027.

• batch every single-contract target in a scope into one Transformer Engine make_graphed_callables() invocation, giving attention and MoE one shared private pool each instead of one pool per layer
• preserve forward layer order; TE captures the corresponding backward graphs in reverse order
• stage explicit FSDP parameter inputs one owner at a time and reshard before TE capture; the fresh capture leaf is also TE's static input, so there is no additional TE-owned parameter copy
• fail closed when one entry in a shared pool no longer matches its captured input/control/parameter contract
• make capture transactional: reset already-created graphs if a later scope fails or TE returns malformed results
• log exact per-entry static-input/parameter-clone bytes and allocator snapshots around staging, plus exact normal/private allocator snapshots around each shared-scope capture

A private-pool delta cannot be attributed honestly to one entry once TE captures the full scope in a single call, so actual graph-pool deltas are reported per scope rather than repeated for every layer.

H100 experiment

Qwen3-30B-A3B, 8x H100 80GB, BF16, FSDP2 + EP8, TE attention and experts, HybridEP, no AC/PP, mock dataset.

Exact old/new comparison at local/global batch 1/8, sequence length 1024:

Graph pool layout	Peak allocated / GPU	Avg step	Result
per-entry pools	53.53 GiB	304.39 ms	pass
shared attention/MoE pools	52.56 GiB	314.83 ms	pass

Shared pools reduce peak allocated memory by 0.97 GiB/GPU (1.8%) at this shape. Runtime is statistically inconclusive in these short runs; this PR's measured benefit is memory/capacity rather than a claimed speedup.

Capacity comparison at local/global batch 2/16, sequence length 1024:

Mode	Peak allocated / GPU	Avg step	Result
eager	39.00 GiB	474.43 ms	pass
old per-entry attention + MoE graphs	>71.32 GiB allocated at failure	n/a	OOM while capturing layer 45/48
shared attention + MoE pools	59.49 GiB	355.56 ms	pass

At this larger shape, shared capture reduces allocated memory by at least 11.83 GiB/GPU relative to the old failure point and turns the previous OOM into successful capture/replay. The completed graph run is 25.1% faster than eager. Eager and graph loss sequences match to the logged precision across all 15 steps.

Allocator diagnostics after shared capture at local batch 2:

• attention shared pool: 5,177.57 MiB private allocated
• attention + MoE shared pools: 21,194.05 MiB private allocated

The larger local-batch-2, sequence-2048 case still OOMs during the MoE backward warmup. That result identifies retained TE expert backward state as the next memory target; a separate stacked PR will page only marked expert saved activations.

Validation

• 170 passed, 9 skipped:
  • tests/unit_tests/recipes/test_partial_cuda_graphs.py
  • tests/unit_tests/recipes/test_train_ft.py
  • tests/unit_tests/recipes/test_train_ft_partial_cuda_graphs.py
  • tests/unit_tests/moe/test_backend_config.py
• Ruff and git diff --check pass
• independent review covered TE ordering, pooled fail-closed behavior, FSDP storage lifetime, failure cleanup, and memory attribution

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.