Pipeline the layer-sharded Muon all_to_all to hide exchange latency (and cut peak memory)

[wanyingw]

Summary

Layer-sharded Muon (LSH) currently performs one all_to_all covering every matrix a
DP shard owns, then runs Newton-Schulz on all matrices homed to this rank, then one
reverse all_to_all. This serializes communication and compute, and holds every
intermediate buffer live simultaneously.

Splitting the exchange into C chunks and overlapping chunk k's Newton-Schulz with
chunk k+1's all_to_all would recover the ~17% of LSH step time currently spent stalled
in NCCL, and gives a memory knob that would otherwise require falling back to
duplicated.

This affects the expert (EGTP) axis, where a DP shard owns 23-24 matrices. Dense (GTP)
profiles mostly own a single matrix, so there is nothing to chunk there.

Current behaviour

layer_sharded_all_to_all_fwd / _bwd (megatron/core/optimizer/layer_sharded_a2a.py)
take the full momentum_list and build one flat send buffer over all of it:

send_parts = []
for g_prime in range(gtp_size):
    if params_for_rank[g_prime]:
        chunk = torch.cat([m.contiguous().flatten() for _, m in params_for_rank[g_prime]])
        send_parts.append(chunk)
send_buf = torch.cat(send_parts) if send_parts else ...
torch.distributed.all_to_all_single(recv_buf, send_buf, ...)

Linked PR

Layer sharded Muon (LSH): #6683