fix(distributed): synchronize tensor-parallel replicas

[yuhezhang-ai]

What does this PR do?

Synchronize parameters that remain replicated across tensor-parallel ranks, both after model construction/loading and once per optimizer update.

Changelog

• Broadcast replicated parameters and buffers from TP rank 0 after model surgery, while leaving intended DTensor shards untouched.
• Coalesce replicated gradients into bounded buffers and reduce them at the optimizer boundary before scaling, norm calculation, and clipping.
• Mean-reduce redundant full-computation replicas and preserve Hugging Face's sum-reduction semantics for partial-head replicated norms.
• Detect asymmetric gradient presence on every TP rank and fail without creating a mismatched collective order.
• Add real two-rank collective coverage for initialization, accumulation, clipping, optimizer updates, unused/frozen parameters, DTensor shards, and folded expert meshes.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? (No user-facing API or configuration change.)

Test evidence

• ruff format . (804 files unchanged)
• ruff check --fix .
• pytest tests/unit_tests/distributed/test_tp_replicas.py tests/unit_tests/training/test_train_utils.py tests/unit_tests/distributed/test_qwen3_5_tp_and_grad_sync.py -q (48 passed, 3 skipped)
• pytest tests/unit_tests/_transformers/test_infrastructure.py tests/unit_tests/distributed/test_parallelizer.py -q (200 passed)
• pytest tests/unit_tests/moe/test_parallelizer.py -q (91 passed)
• Nano 9B TP=2 PEFT save/reload on 8 GPUs: passed (Slurm 17862116). Trainable-parameter fingerprints matched exactly on all 8 ranks; reloaded logits had cosine similarity 1.0 and zero maximum absolute difference.
• Nano 9B TP=2 full-SFT save/resume on 8 GPUs: passed (Slurm 17862129). Model parameters, buffers, optimizer tensors, optimizer counters/groups, scheduler, and RNG matched at the first common pre-update point on all 8 ranks. The first restarted-forward loss delta was 1.8148e-3, within the 4.1175e-3 tolerance, and decreased on later steps.

Performance

In these short four-step 8-GPU probes, PEFT averaged 27.70 s/step versus 27.57 s/step before the fix (+0.5%), while full SFT averaged 35.00 s/step versus 34.27 s/step (+2.1%). The fix coalesces replicated gradients into bounded buffers, but full SFT necessarily communicates the parameters that the current TP plan leaves replicated.

Additional Information

• Fixes the underlying TP replica initialization problem reported in Custom module weights not synchronized across TP ranks #1291.
• Covers the checkpoint-state portion of AM-212, AMINT-291, and AMINT-307. AMINT-314 and the separate restarted-forward numerical drift in AMINT-307 are intentionally out of scope.
• Supersedes the replica synchronization scope of fix: tp weight sync #1292. Credit to @akoumpa for the earlier investigation and initialization-broadcast implementation.
• This is a TP state-consistency fix, not a checkpoint-loader change: the loader was faithfully restoring one of several already-diverged replicas.

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