[Dev] Make FSDP manual UBR registration deterministic

[lhb8125]

Summary

• pre-materialize every runtime-reachable persistent FSDP/MaxPool buffer before manual NCCL user-buffer registration
• build and validate a globally stable logical allocation plan across every registration communicator
• validate the physical memory-pool segment signature before entering collective register_mem_pool, so a mismatch fails with diagnostics instead of hanging peers
• use exact padded runtime bucket sizes and canonical buffer names rather than relying on rank-local first-touch order from overlapping streams

Symmetric registration requires matching collective registration order and pool-relative layouts across ranks. Lazy allocation during the first combined-1F1B/CUDA-Graph iteration allowed identical logical requests to receive different physical segment orderings on different ranks.

Stack

This Draft PR is stacked on #6316. After #6316 merges, this branch will be rebased onto dev; the remaining diff is the single fix(fsdp): make manual UBR registration deterministic commit.

Test plan

• added unit coverage for registration-counter ordering, rank-local-address exclusion, and exact MaxPool materialization requests
• CHECK_ONLY=true BASE_REF=dev ... tools/autoformat.sh: Black, isort, pylint, and ruff pass
• copyright check passes for all changed Python files
• large-scale reproduction: job 2592505 exposed a rank-dependent 117 MB physical segment; the materialization/global-plan sequence advanced the same 61-layer full-CG proxy to a deterministic pool layout in 2592639

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.