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.