Fix ZeRO-3 Muon all-gather for mixed parameter shapes

[jinyouzhi]

Follow up #8381
Fix 3 FAILED unit/v1/ops/muon/test_per_head_muon_accelerator.py in https://github.com/deepspeedai/DeepSpeed/actions/runs/35675316100/job/106580600974?pr=8381
@delock

[alanhuangyoo]

Thanks for picking this up! I think this needs a rebase on #8600, the grads come from group_items there now since param.grad isn't the full grad at step time.

Btw NCCL seems fine with uneven all_gather, which is probably why only the gloo job failed.

[jinyouzhi]

Thanks for picking this up! I think this needs a rebase on #8600, the grads come from group_items there now since param.grad isn't the full grad at step time.

Btw NCCL seems fine with uneven all_gather, which is probably why only the gloo job failed.

Thank you for remind! Yes, the ut only failed in Gloo, NCCL better compatible on this case.

[alanhuangyoo]

Thanks, the rebased version looks good to me. The ZeRO-3 Muon tests pass on 2×H20, and ZeRO-3 still matches ZeRO-0 bit for bit with mixed shapes.

[jinyouzhi]

Thanks! @alanhuangyoo
@delock PTAL