[Dev] Scope FSDP symmetric UBR to dense inner gathers

[lhb8125]

Summary

• add an opt-in fsdp_ubr_registration_scope=dense_inner mode that registers only the dense inner-FSDP parameter-AllGather communicator
• keep expert, outer-DP, gradient, and all-in-one storage on ordinary allocation/collective paths
• isolate dense and expert MaxPool allocators so expert first-touch order cannot perturb the registered pool
• pack every registered dense parameter/transpose/persistent allocation into one deterministic 256-byte-aligned arena
• reject unsupported configurations early and preserve the existing all registration scope as the default

The scoped mode targets HSDP workloads where dense inner AG dominates overlap pressure while expert outer-DP communication is small. Rank-offset HSDP ReduceScatter outputs are intentionally excluded because the current symmetric-window path requires identical pool-relative offsets across ranks.

This PR deliberately does not include fsdp_ubr_arena_segment_bytes: 5.8 GB, 4 GiB, and 2 GiB registration attempts all failed on the same unintended SPCX GIN path, while restoring the intended built-in GDAKI path made the unsegmented design work. Arena segmentation is therefore not a demonstrated Megatron requirement.

Stack

This Draft PR is stacked on #6317, which is stacked on #6316. It should be reviewed as the single perf(fsdp): scope symmetric UBR to dense inner gathers commit after the prerequisite PRs merge and the branch is rebased onto dev.

Test plan

• unit coverage for communicator selection, unknown-scope rejection, dense/expert pool isolation, aligned arena layout, and deterministic materialization
• CHECK_ONLY=true BASE_REF=dev ... tools/autoformat.sh: Black, isort, pylint, and ruff pass
• copyright check passes for all changed Python files
• historical validation: scope-only job 2585550 and pool-isolation job 2585577 reproduced remaining ordering failures; packed-arena proxy 2585621 and 256-GPU full-model job 2585683 completed
• latest-dev validation: 64-GPU, 61-layer full-CG job 2593409 completed 11/11 iterations with NCCL 2.31 MR3058; dense-inner AG used ncclSymkDevKernel_AllGather_RailRing_LsaSTMC, reduced the matching ordinary-ring grid from 16 to 8, and reduced aggregate dense AG kernel time by 31.36%

The local toolkit lint environment does not contain PyTorch, so pytest and mypy imports are deferred to CI/container validation.

🤖 Generated with Codex

[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.