[GTP][Muon] Optimize Muon step performance

[wanyingw]

Project Overview

• Goal: Build the underlying architecture for layer-sharded distributed Muon optimization in large-scale tensor-parallel training. Each weight matrix's Newton-Schulz orthogonalization is assigned to a dedicated home rank in the parallel domain (all_to_all gather → local full-matrix compute → scatter back), eliminating the redundant compute and per-parameter collectives of the existing duplicated mode.
• Owner: @wanyingw

Current Progress

Primary workstream — Layer-sharded Muon

• Architecture & design (layer-sharding strategy, two-stage → fused single all_to_all, LPT load balancing)
  Core implementation (confidential repos): optimizer + training-framework integration (process groups / CLI / dense & expert dual-domain routing)
• Correctness verification: 28 unit tests (CPU & GPU backends), optimizer-level bitwise parity vs reference, real-model end-to-end loss comparison with a run-to-run-noise control experiment
• Performance: shape-batched NS, fused all_to_all, SYRK symmetric kernels (validation in progress)
• Backward/forward overlap with dynamic wave dispatch (next milestone): NS work is statically partitioned into waves sized to the GPU allocation (e.g. 3000 tensors / 1000 slots → 3 waves; slots > tensors → surplus goes to duplicated homes for scatter locality or team-NS for latency), triggered dynamically as gradient buckets complete — gradient gather hidden in the backward shadow, NS + scatter hidden in the next iteration's forward
• At-scale performance profiling (real-model breakdown vs duplicated baseline)
  Feature completion: split-QKV support, MoE end-to-end validation
  Integration readiness / upstream merge review

Secondary (parked) — early explorations

• Distributed NS with single-all-reduce Gram (gram_once): exact Gram-space reformulation, 8× fewer collectives, bf16-instability root-caused and hardened (fp32/TF32 Gram + trace ridge). Parked due to numerical instability.
• Refresh-4: similar to gram-once reformulation, with 1/4 the all-reduce needed in distributed gram, full fp32 preserves numerical accuracy, cluster-validated. Net win is topology-dependent (targets inter-node regimes).Parked in favor of layer sharding.

Note

Development happens in separate repositories (optimizer library + training framework) and will be merged upstream once stabilized/cleared. The work does not change the optimizer's mathematical semantics, only restructures how the computation is distributed, backed by bitwise-level verification. Please reach out to me directly for details.