[Dev] Generalize FSDP persistent communication buffers

[lhb8125]

Summary

• generalize Megatron-FSDP's persistent communication pools from a hard-coded double buffer to a configurable fsdp_buffer_count
• require at least three persistent slots when combined 1F1B overlap is used with FSDP persistent buffers
• account for buffers that are already live, as well as buffers that can be released before the next allocation, when deciding whether parameter prefetch fits
• keep the default at two slots for backward compatibility and expose the setting through Megatron-Core and standalone fully_shard APIs

Combined 1F1B can keep a backward/recompute unit, the current genuine-forward unit, and its forward-prefetched successor live concurrently. Counting only requested FSDP units can overwrite a live allocation or reject a valid prefetch because different dtype/offset pools have different availability.

Test plan

• added CPU unit coverage for FixedPoolAllocator and MaxPoolAllocator capacity checks with live and releasable buffers
• CHECK_ONLY=true BASE_REF=dev ... tools/autoformat.sh: Black, isort, pylint, and ruff pass
• copyright check passes for all changed Python files
• validated as part of the 61-layer DeepSeek-V3 combined-1F1B/full-iteration-CUDA-Graph runs on GB200: SLURM 2593409 completed 11/11 iterations; earlier two-slot and live-accounting failures reproduced in 2592350 and 2592413

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.