[Dev] Add zero-CTA all-gather policy for Megatron-FSDP

[bobboli]

Summary

• Add --megatron-fsdp-zero-sm-all-gather to request NCCL_CTA_POLICY_ZERO on dedicated dense and expert Megatron-FSDP parameter all-gather communicators.
• Preserve inherited NCCL communicator settings while allowing per-all-gather overrides, and validate the required Megatron-FSDP, user-buffer, symmetric-registration, and process-group configuration.
• Keep the policy best-effort: eligible registered buffers use copy engines; ineligible buffers retain NCCL's kernel path without changing FSDP allocator fallback behavior.
• Treat an installed NVRX version below 0.6.0 as unavailable during optional capability detection; explicit NVRX selection still raises a clear error. This compatibility fix was required by the available OCI image and is kept as a separate commit.
• Document CUDA/NCCL version requirements and the latency-versus-SM-availability tradeoff.

Functional validation

On one 4x GB200 OCI node with PyTorch 2.12.0a0 and NCCL 2.29.7, an order-balanced 1 MiB/rank probe ran default, zero, zero, default with 2,000 timed calls per run:

Policy	Mean latency	Profiler evidence
Default	0.03736 ms	Five ncclSymkDevKernel_AllGather_STMC kernels for five profiled calls
Zero	0.08338 ms	No all-gather SM kernel; copy-engine P2P/D2D transfers only

All runs were numerically correct and the communicator reported policy values 0 and 2 respectively. The isolated zero-CTA collective was 2.23x slower, which matches NCCL's documented policy tradeoff: it preserves CTAs for compute and may sacrifice communication latency.

Qwen3.5 performance

Qwen3.5-35B proxy, Megatron-FSDP EP4, four GB200 GPUs, sequence length 4096, 20 iterations, eight warmup iterations skipped:

• Same node, baseline then zero: 2901.6 ms versus 2849.9 ms median (1.78% faster).
• Same node, reverse order: zero 2971.15 ms versus baseline 2957.0 ms (0.48% slower).
• Peak allocated/reserved/device-used memory was identical.
• Matched traces showed parameter all-gather time increasing from 5.947 ms to 6.949 ms across 26 calls (+16.8%). The apparent whole-step speedup came from unrelated GPU-idle variance.
• The Qwen diagnostic logged 72 dynamic-allocation fallbacks where NCCL user buffers were unsupported, and baseline/candidate traces had identical all-gather GPU event counts. This workload therefore did not exercise zero CTA for the measured FSDP all-gathers.

There is no end-to-end throughput claim in this PR. The code exposes a verified zero-CTA path for eligible buffers; heterogeneous Qwen FSDP buffer eligibility remains follow-up allocator work.

Test plan

• Focused CI-image suite: 20 passed, 227 deselected.
• Black, isort, pylint, ruff, and copyright checks passed.
• Single-node controlled zero-CTA probe passed on NCCL 2.29.7.
• Cross-network zero CTA was not tested because NCCL requires 2.30.6+ and the available OCI image contains 2.29.7.

References: NCCL zero-CTA requirements, NCCL CTA policy semantics.

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