Fix ZeRO-3 Muon all-gather for mixed parameter shapes (#8628)
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
---------
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: pengdurice <pengduhit@gmail.com>
