DeepSpeed
e94a8da1 - fix(autoep+muon): auto-fill HF config fields and restore batched NS

Commit
126 days ago
fix(autoep+muon): auto-fill HF config fields and restore batched NS Two fixes addressing masahiro's review feedback on PR #7938: 1. Auto-fill AutoEPConfig from HF model config (auto_ep_config.py, auto_ep.py): add fill_autoep_config_from_hf() which maps HF field names to AutoEP internal names on AutoEP.__init__: - n_group -> num_expert_groups - topk_group -> num_limited_groups - routed_scaling_factor -> route_scale User-supplied values always take precedence. Without this, Moonlight (DeepSeek-V3) training used route_scale=1.0 instead of 2.446, producing systematically wrong MoE output magnitudes. 2. Restore batched Newton-Schulz in muon_update (original_muon.py): replace the per-expert Python loop with a single batched call to zeropower_via_newtonschulz5, which already supports ndim>=2 inputs. This restores GPU parallelism across all E experts per step. Signed-off-by: Ma, Guokai <guokai.ma@gmail.com>
Author
Committer
Parents
Loading