fix(muon): support ZeRO-1/2 reduce scatter (#8141)
## Problem
#8090 made `Muon + ZeRO-1/2 + reduce_scatter` fail fast because a Muon
matrix that crosses a ZeRO partition boundary does not have a fully
reduced gradient on any owner rank.
The reduce-scatter path copies each averaged slice only to that slice's
owner. Muon, however, runs the nonlinear Newton-Schulz update on the
full matrix before `get_flat_partition()` selects the local ZeRO slice.
Each owner therefore sees a different matrix: its own reduced slice plus
local, unreduced slices owned by other ranks.
## Change
For a Muon parameter that spans more than one ZeRO partition, each
reduced slice is copied to every rank that owns part of that parameter.
Those ranks then run `muon_update` on the same full reduced gradient,
and the existing code keeps only the local ZeRO partition.
Parameters that stay within one partition keep the current owner-only
copy path. With the default multi-rank bucket all-reduce, this does not
add a collective. When `use_multi_rank_bucket_allreduce=false`, only
ranges belonging to split Muon parameters switch to all-reduce. The copy
helper still accepts the integer rank form used by ZenFlow.
Optimizer offload remains unsupported because that path keeps only
partition-local gradients.
## Tests
The numerical test compares the applied update with `muon_update` on the
full data-parallel gradient. The 2-GPU cases cover:
- ZeRO stages 1 and 2
- `gram` and `standard` Newton-Schulz methods on the existing all-reduce
path
- reduce-scatter with gradient accumulation and overlap communication
- `use_multi_rank_bucket_allreduce=false`
- extra-large parameters and non-contiguous gradients
- optimizer offload rejection
Focused helper tests cover both the existing integer copy target and the
new multi-owner target.
Refs #7807
Alternative to #7878; replaces the fail-fast behavior added by #8090
with a targeted reduction path.
Signed-off-by: wuqingliu <wuqingliu@kuaishou.com>
Co-authored-by: wuqingliu <wuqingliu@kuaishou.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
