Fix megabatch_orthogonalize_async hang on empty FSDP2 shards

[JohnLangford]

Summary

megabatch_orthogonalize_async hangs at NCCL dist.all_to_all whenever a managed 2D weight's sharded global dim is smaller than world_size or doesn't divide evenly to fill all ranks under FSDP2's contiguous chunking. Some ranks receive empty (0, …) padding-only shards; the per-rank input_chunks then have inconsistent sizes on the wire and the alltoall deadlocks after the first backward.

Concrete repro (the case that surfaced this)

• 2D weight of shape (num_attention_heads=18, hidden=2304) (a sparse-3B model's per-head attention gate).
• FSDP2 fully_shard over dp_shard_size=8.
• NorMuon (or Dion2) — both hit the same path.

_chunk_with_empty(t, 8, dim=0) returns 6 real (3, 2304) chunks plus 2 empty (0, 2304) chunks. After init: ranks 0–5 hold (3, 2304) shards; ranks 6–7 hold (0, 2304). Inside megabatch_orthogonalize_async:

```python
input_chunks = [torch.stack(U_work[r * per_rank : (r + 1) * per_rank]) for r in range(world_size)]
dist.all_to_all(output_chunks, input_chunks, group=process_group)
```

Each rank constructs input_chunks from its own local data, so per-target chunk shapes diverge across ranks ((per_rank, 3, 2304) vs (per_rank, 0, 2304)). dist.all_to_all requires per-pair sizes to match — they don't.

NCCL watchdog signature for 32 same-shape params packed under per_rank=4:

```
Ranks 0–5: WorkNCCL(SeqNum=N, ALLTOALL, NumelIn=8×4×3×2304 = 221184)
Ranks 6–7: WorkNCCL(SeqNum=N+1, ALLTOALL, NumelIn=8×7×96×2304 = 12386304)
```

The 6/2 watchdog split exactly matches torch.chunk(18, 8, dim=0). Padding-only ranks' zero-sized alltoall completes locally so they advance to the next pending collective and hang there.

Independent of optimizer (NorMuon and Dion2 both hit it) and of any upstream-loss specifics; it's a structural property of the megabatch sharded path.

Fix

Pad each rank's local shard along comm_dim to a rank-consistent padded_local_size = max(local_sizes) before the alltoall (one extra dist.all_reduce MAX), then narrow the post-alltoall result back to the rank's original local size. Newton-Schulz preserves zero rows (they contribute nothing to UᵀU), so padding doesn't change the orthogonalization of the real rows. Mirrors what FSDP2 itself does for its own all-gather/reduce-scatter (padded_sharded_param_size in torch/distributed/fsdp/_fully_shard/_fsdp_param.py).

Test plan

• New test tests/test_megabatch_empty_shard.py spawns 4 NCCL ranks via torch.multiprocessing.spawn and exercises three configurations:
  • (global_dim=5, world_size=4) — rank 3 holds (0, …); analogous to the (18, 8) production case at smaller scale.
  • (global_dim=2, world_size=4) — ranks 2 and 3 both hold (0, …).
  • (global_dim=15, world_size=4) — non-divisible but no empty ranks (sanity check that the padding logic stays a no-op when uniform).
• Verified the test hangs (and the timeout fires) on main without the fix and passes with it. All 3 cases pass with the fix; full existing suite (97 tests) continues to pass.
• End-to-end confirmation on a real cluster: a phitrain sparse-3B run that previously hung on the first optimizer step now trains 39 steps cleanly with loss decreasing smoothly (6.95 → 6.45).

[JohnLangford]

@alint77 can you check this one?

[alint77]

LGTM.

[alint77]

this has introduced a sync point in the pipeline, im working on a fix atm, will submit PR soon

[JohnLangford]

Ah? Yes please!