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: CPUOffloadPool.reload_group에 sync_streams 파라미터 추가 (#31)
Unverified signature* feat: tag-based per-group reload for CPU offload + manual offload mode Extend CPUOffloadPool + Muon so callers can drive offload/reload explicitly (layer-lockstep overlap) rather than implicitly inside optimizer.step. CPUOffloadPool (torch-ext/optimizer/cpu_offload.py): * ``track(tensor, tag=None)`` — optional tag per managed tensor. * ``reload_group(tag, sync_streams=())`` — reload just the tensors tagged with ``tag``; the reload stream ``wait_stream`` s on each entry in ``sync_streams`` before issuing H2D. This avoids allocator cross-stream reuse races under ``PYTORCH_ALLOC_CONF=expandable_segments:True``: if the block returned by ``storage.resize_`` was last touched on an FSDP all-gather stream, waiting on that stream enforces FIFO ordering between the prior use and our H2D write. * ``reload_untagged()`` — bulk-reload everything not attached to a tag (for the non-expert portion of the optimizer state in layer- lockstep flows). * ``wait_reload()`` is now self-c…d7e2d58
-
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
-
Replace cpu_offload constructor param with turn_on/turn_off API (#26)
Unverified signature* Replace cpu_offload constructor param with turn_on/turn_off API [skip-build] Remove the cpu_offload boolean from Muon.__init__ and add explicit turn_on_cpu_offload() / turn_off_cpu_offload() methods instead. state_dict and load_state_dict now require offload to be disabled first (RuntimeError). Preserves AdamW tensor cache invalidation on load_state_dict. Adds test_toggle_correctness test. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: skip empty FSDP shards in CPU offload and add storage validation Empty FSDP shards (storage size=0) were being tracked by CPUOffloadPool, causing double-free errors on resize_(0) during offload. This led to hangs on ranks 1-7 while rank 0 succeeded. - Skip tensors with empty storage in track() to avoid double-free - Add storage size validation in offload() and reload() with RuntimeError - Add logging for turn_on/turn_off_cpu_offload Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Add built bina…05a75f1
-
draft commit for cpu_offload (#23)
TaehyunKim +3: github-actions[bot], wyldecat, Claude Opus 4.6 (1M context)Unverified signature* draft commit for cpu_offload * draft commit for cpu_offload [skip-build] * claude is british * Add built binary [skip-build] * fix load, save and add test * Add built binary [skip-build] * Fix yapf formatting for pre-commit CI [skip-build] Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: WyldeCat <skan1543@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>10848ab
-
Replace toy PP tests with real-model-based pipeline tests [skip-build]
Use Motif-2.6B (dense) and torchtitan Llama4 MoE (MoE) models with realistic PP model splitting (deep copy → delete non-stage layers → per-stage FSDP) matching actual torchtitan training pipeline. MoE test uses torchtitan's parallelize_llama() directly. Both tests verify correctness against sequential baseline with atol=0, rtol=0. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>67f7e11 -
Add correctness verification to PP tests using fully_shard [skip-build]
Use fully_shard (proven HSDP pattern) instead of manual distribute_tensor to create proper DTensors. Verify parallel results match sequential baseline with atol=0, rtol=0. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>a4d1f34 -
Remove correctness check from PP tests, focus on deadlock detection [skip-build]
Correctness of Replicate+Shard is already verified by existing hsdp tests. These PP tests specifically validate that asymmetric construct_shard_mesh calls across PP stages don't deadlock. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>c0bbf2e -
Add PP + dp_replicate deadlock regression tests [skip-build]
Add tests simulating Pipeline Parallelism with dp_replicate > 1 to verify construct_shard_mesh doesn't deadlock when different PP stages call dist.new_group for different parameters independently. - test_muon.py: Dense model test (PP=2, dp_replicate=2, dp_shard=2) - test_muon_moe.py: MoE test with asymmetric expert/non-expert params - test_pp.yaml: K8s job spec for running on GPU pod Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>cd587a6
-
Apply pre-commit formatting (isort) [skip-build]
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>96b287c
-
Add MoE uneven shard test with mixed expert and non-expert params [skip-build]
Test parallel Muon with uneven dims (33, 19) mixing 2D DTensor params (parallel pipeline) and 3D expert plain tensors (batched NS path). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>bdada12 -
Add uneven shard correctness test [skip-build]
Test parallel Muon with param dimensions not divisible by shard count (dim=33,19,11 with 8 ranks). Verifies against sequential baseline. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>1a97671 -
Update tests for MoE and parallel optimizations [skip-build]
- Add MoE test cases (test_muon_moe.py) - Update parallel test configurations - Test utility updates Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>81f49fe -
Add torch.compile, CUDA graph, and compiled momentum [skip-build]
- Newton-Schulz: per-shape torch.compile caching + CUDA graph replay - Batched momentum: separately compiled nesterov/non-nesterov functions - Batched Newton-Schulz for MoE experts (bmm/baddbmm) - Triton matmul_transpose cleanup - Inline uneven shard handling, remove small_param_numel_threshold - Raise dynamo recompile_limit for test suite Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>e74d98f
-
Apply suggestions from code review
TaehyunKim +1: CopilotCo-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>cdaaf4f
-
Support multi-component expert_keys (e.g. "experts.w1")
Single-component keys match any single FQN component as before. Multi-component keys (containing dots) now match as a contiguous subsequence, so "experts.w1" matches "moe.experts.w1.weight" but not "moe.shared_experts.w1.weight". [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>5a99e12 -
Use component-level matching for expert_keys to avoid shared_experts collision
Substring matching (`key in name`) causes "experts" to match "shared_experts". Switch to dot-split component exact matching (`key in name.split(".")`) in default_is_muon() and _expand_expert_params(). Also applies to skip_keys for consistency. [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>f008017 -
Normalize parameter FQNs to handle torch.compile / checkpoint wrappers
torch.compile wraps modules as OptimizedModule, inserting _orig_mod into parameter FQNs. Activation checkpointing similarly inserts _checkpoint_wrapped_module. These wrapper components break name-based matching for skip_keys, expert_keys, and QK layer parsing. Add normalize_fqn() that strips these wrapper components, and apply it in default_is_muon(), _expand_expert_params(), and parse_qk_layer(). [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>95a620f
-
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 mHC (#15)
Unverified signature* log target shape * update toml * use distributed muon for small param * Add built binary [skip-build] --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>ae32572
-
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
-
misc
35894d1 -
apply all2all scatter gather
ff6d675 -
feat(muon_clip) : add muon clip (#6)
Unverified signature* feat(muon_clip) : add muon clip * fix(muon_clip): delete comment * fix(muon_clip): delete comment * fix(muon_clip): considering when nkvheadgroup>1 * docs(muon_clip): refine __init__ docstring and add clip_info argument description * refactor(muon_clip): refactor clip info using dataclass * fix(muon_clip): change min -> new_scaling compare * test(muon): add qk_clip=False case to model comparison * test(muon): show results * fix(muon_clip): change default is muon func * Add built binary [ci skip] --------- Co-authored-by: dongseokmotif <dongseokkim@motiftech.io> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>d65066c
-
feat: update muon to receive paramgroups, not model (#4)
Unverified signature* feat: update muon to receive paramgroups, not model * feat: update message formats * feat: remove boilerplate assertion * chore: run pre-commits * test: fix testscript * fix: fix misc bugs * feat: add get_default_muon_param_groups helper function * fix: fix readme * fix: raise error if parametergroup does not follow instructions * chore: upload binary --------- Co-authored-by: junhyeok.lee <junhyeok.lee@moreh.io> Co-authored-by: WyldeCat <skan1543@gmail.com>b0f46c7
-
feat(muon): add test for muon
3261444