Commit History
-
feat: add Per-Head Muon (Kimi K3 §2.5) via head_keys [skip-build]
Partition attention Q/K/V momentum matrices along the head dimension (dim 0) and orthogonalize each head's block separately, instead of running Newton-Schulz on the full projection matrix. Equalizes the update scale across heads (K3 tech report §2.5). - head_keys={"q_proj": 128, ...}: component-level name matching (same rule as expert_keys) mapping to per-head dim-0 row stride; MLA strides supported (e.g. wq_b/wkv_b) - NS granularity only: momentum stays elementwise; matching params are viewed [H*d, N] -> [H, d, N] and run through the existing batched bmm Newton-Schulz; adjusted LR uses the head-block shape (consistent with the per-expert path) - Sharding-agnostic: all three paths (base / distributed_muon / parallel) apply the head partition to the full gathered matrix, so results are identical under any DP/TP/HSDP placement with no communication changes; _muon_state.head_dim switches the compute call - conftest: make transformers import lazy so synthetic tests run in…ebeb3a1
-
feat: extend QK-Clip to support MLA (MuonClip Algorithm 1) [skip-build] (#28)
Unverified signature* feat: extend QK-Clip to support MLA (MuonClip Algorithm 1) [skip-build] Add MLA support to qk_clip.py following the MuonClip spec (docs/muon-clip.md): - parse_qk_layer: recognize 'wq_b' and 'wkv_b' MLA weight names - QKClipInfo: add is_mla, qk_nope_head_dim, qk_rope_head_dim, v_head_dim fields - get_qk_clip_info: branch on is_mla flag in clip_config - compute_scales: use kv_stride (qk_nope + v_head_dim) as effective head dim for wkv_b - qk_clip: simplify signature to (p, scales, info); vectorize MLA sub-region scaling via tensor reshape instead of Python per-head loops: wq_b: q_nope rows → √γ, q_pe rows → γ wkv_b: k_nope rows → √γ, v rows unchanged Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: ruff format (pre-commit fix) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Add built binary [skip-build] --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: WyldeCat <skan1543@gmail.com> Co-authored-by: g…e8e2c81
-
Add optimization docs and update implementation guide [skip-build]
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>14040eb
-
Refactor pipeline to async generator pattern (#16)
Unverified signature* Refactor muon.py into modules with async generator pipeline * Add MoE expert weight support with EP+FSDP tests * Add built binary [skip-build] --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>33929c0
-
Support param group with various placements (#13)
Unverified signature* feat(muon): group parameters by placements for parallel Muon execution * refactor(muon): refactor step func and group params with it's placement * feat(muon): support general mesh * refactor(muon): refactor state init * refactor(muon): refactor test * fix(muon): fix general mesh, add chunk_size argument * refactor(muon): change overlap_step to warmup_step * refactor(muon-test): rewrite README, add conftest.py and use explicit flags * chore(muon): clarify N-D sharding support and add test reference * fix: use device_mesh as key to group params * Add built binary [skip-build] --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>e2b41e5
-
chore(muon): clean build and update doc
febdf5b
-
chore: initial commit
8535e80