Commit History
17 days ago
-
Merge pull request #33 from MotifTechnologies/feat/per-head-muon
Unverified signaturefeat: add Per-Head Muon (Kimi K3 §2.5) via head_keys [skip-build]a85f310
-
Merge pull request #34 from MotifTechnologies/fix/qk-clip-rotary-full-gamma
Jangwoong KimUnverified signaturefix(qk_clip): q_pe(rotary query)를 full γ로 스케일 (γ^(1-k_ratio) 아님)9765db5
-
qk_clip: value-only fix — keep original structure, drop per-row rewrite [skip-build]
Revert the per-row expand_row_scales restructure. qk_clip keeps its exact original reshape + named-slice shape; only the multiplied values change: compute_scales now returns raw per-head γ, and each sub-region is raised to its exponent inline (q_nope γ^(1-r), q_pe γ, k_nope γ^r, MHA √γ). The q_pe line is byte-identical to before — it is now correct because scales is raw γ. muon.py sharded path stays per-head (no per-row vectors); it just applies the content exponent to match the raw-γ contract, preserving its prior behavior. The sharded path's pre-existing MLA sub-region limitation (q_pe/v scaled by the content exponent) is left as-is and documented, not addressed here. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>0669c49 -
refactor(qk_clip): fold row-scale logic into one expand_row_scales helper [skip-build]
Collapse _row_exponents + expand_row_scales into a single expand_row_scales using cat + broadcast (gamma[:,None] ** exps[None,:]).reshape(-1) instead of repeat_interleave/repeat. qk_clip stays two lines; muon.py unchanged. No behavior change — pure readability. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>12c06f0 -
refactor(qk_clip): apply γ exponent per-row so q_pe→full γ in both paths [skip-build]
Replace the round-trip fix with a clean design: compute_scales now returns the raw per-head γ (= threshold/value) and the per-region exponent on γ is applied at multiply time by a shared expand_row_scales() helper. Row exponents: MHA : √γ on all rows MLA q_nope : γ^(1-k_ratio) ┐ content = γ^(1-k)·γ^k = γ MLA k_nope : γ^k_ratio ┘ MLA q_pe : γ (exponent 1) ─ rotary = γ·1 = γ (k_rope shared/untouched) MLA v rows : γ^0 = 1 ─ untouched This removes the ugly scales**(1/(1-k_ratio)) round-trip and its broken k_ratio=1 branch: since γ is raw, q_pe just uses exponent 1 directly and every k_ratio in [0,1] is correct with no division. It also fixes the sharded DTensor path in muon.py, which bypassed qk_clip and applied one per-head scale to all rows — that both under-clipped q_pe and wrongly scaled the wkv_b v rows. Both apply paths now go through expand_row_scales (qk_clip multiplies the whole tensor; the DTensor path indexes the per-row vector by this …3b17f3b -
fix(qk_clip): scale q_pe rows by full γ, not γ^(1-k_ratio) [skip-build]
MLA logit = q_nope·k_nope (content) + q_pe·k_rope (rotary). The decoupled rotary key k_rope is shared across heads and left unchanged, so the rotary term is reduced by the query side alone. Since PR #32 introduced k_ratio, wq_b scaled q_pe by γ^(1-k_ratio) uniformly with q_nope, making the rotary term shrink by γ^(1-k_ratio) instead of γ — an under-clip for any k_ratio > 0. Give q_pe the full γ (exponent 1) regardless of k_ratio; the k_ratio split only applies to the content (nope) rows, whose K side is scalable. k_ratio=0 is unaffected (Q already got full γ there). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>349b794 -
Merge pull request #32 from MotifTechnologies/feat/qk-clip-k-ratio
Jangwoong KimUnverified signaturefeat: add k_ratio for Q/K split control in MLA QK clipping [skip-build]8d60c30
-
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: add k_ratio for Q/K split control in MLA QK clipping [skip-build]
dongseokmotif +1: Claude Sonnet 4.6- QKClipInfo에 k_ratio 필드 추가 (기본값 0.0 = 현재 동작 유지) - compute_scales: wkv_b를 k_ratio>0일 때 활성화, exponent를 γ^(1-r)/γ^r로 분리. 근거: K gradient의 유효 신호가 √G 스케일로 감소하므로 clip/recovery 균등 조건 → r = 1/(1+√G) - wkv_b H_global 계산 버그 수정: head_dim 대신 qk_nope+v_head_dim 사용 - compute_scales에 name 파라미터 추가 (로그 식별용) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>9e14f5a
-
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 -
Revert "fix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibility" (#29)
Unverified signature* Revert "fix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibility" * Add built binary [skip-build] --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>313d56a
-
Merge pull request #27 from MotifTechnologies/fix/cudagraph-offload-conflict
TaehyunKimUnverified signaturefix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibilitydebafde -
Add built binary [skip-build]
github-actions[bot]e2f5aab -
fix: disable CUDA graphs in Newton-Schulz for cpu_offload compatibility
CUDA graphs bake in memory addresses at capture time. When cpu_offload frees optimizer state storage via resize_(0) between steps, the CUDA graph replay accesses stale addresses, causing illegal memory access. Remove triton.cudagraphs option and cudagraph_mark_step_begin() from Newton-Schulz compile. torch.compile is preserved for kernel optimization. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>2dce952
-
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
-
Merge pull request #25 from MotifTechnologies/fix/invalidate_cache_adamw
TaehyunKimUnverified signatureInvalidate AdamW tensor caches on load_state_dict [skip-build]b61425a -
Invalidate AdamW tensor caches on load_state_dict [skip-build]
Clear _placement_cache and _tensor_cache after loading state to prevent stale tensor references from causing incorrect optimizer steps. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>89b6099
-
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 -
Update fast path comment to reflect current behavior [skip-build]
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>7e33533 -
Update comment to reflect use_local_synchronization behavior [skip-build]
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>3f5cf49 -
Fix deadlock in construct_shard_mesh with PP + dp_replicate > 1
When Pipeline Parallelism is combined with dp_replicate > 1, different PP stages own different parameters and call dist.new_group() in different orders, causing a collective mismatch deadlock. Fix by using use_local_synchronization=True so only ranks within the same group need to coordinate, and skip creating groups for shard meshes the current rank doesn't belong to. [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>da7e5da
-
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 -
Add optimization docs and update implementation guide [skip-build]
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>14040eb -
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 -
Muon optimizer: expert batching, parallel caching, A2A overlap [skip-build]
- Batched expert NS path for plain-tensor MoE params (skip expansion) - Expert expansion cache to eliminate per-step detach overhead - _setup_parallel() extraction for parallel metadata reuse - Prelaunch first chunk A2A gather to overlap with expert NS compute - Profiler annotations and clarify distributed_muon as test-only Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>0f37d63 -
Optimize pipeline: batched update, zero-copy scatter, prelaunch gather [skip-build]
- Batched update_params with _foreach_mul_/_foreach_add_ - Zero-copy scatter directly into param views - Deferred QK clip buffer allocation - Buffer packing optimization - prelaunch_first_gather() for overlapping first chunk A2A with compute Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>2816b64 -
Cache AdamW placement grouping and tensor lists [skip-build]
- DTensor unwrap via _to_local() to eliminate __torch_dispatch__ overhead - _placement_cache: group params by placement once, reuse every step - _tensor_cache: cache (params_local, moment1, moment2, state_steps) - Only gradients collected fresh each step - Profiler annotations for AdamW stages Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>8ca2492 -
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 -
Add mhc_attn, mhc_ffn, lambda_proj to skip_keys
These params (mHC attention/FFN projections, lambda projections) should use AdamW instead of Muon. [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>ba293d0
-
Remove verbose param_groups summary logging
Per-param logging in default_is_muon() already covers each parameter. The summary list was redundant and too noisy. [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>24f0957 -
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 -
Extract is_expert_param() helper to consolidate expert key matching
The same normalize_fqn + component-level matching logic existed in both default_is_muon() and _expand_expert_params(). Extract into a single is_expert_param() function in core.py so the logic lives in one place. [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>e615b1c -
Include original (pre-normalize) FQN in is_muon logging
Show both normalized and original parameter names so wrapper-injected components (_orig_mod, _checkpoint_wrapped_module) are visible in logs. [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>135fc66 -
Add info-level logging for param group classification (Muon vs AdamW)
Log each parameter's skip/expert status and effective ndim in default_is_muon(), and summarize Muon/AdamW param lists in get_default_muon_param_groups(). [skip-build] Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>1118752 -
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
-
Merge pull request #17 from MotifTechnologies/optimal-ns-coefficients
Unverified signatureReplace hardcoded NS coefficients with analytically optimal ones [ski…b220459 -
Apply pre-commit formatting (yapf) [skip-build]
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>bf30b9b -
Add max_iter cap and non-finite checks to _optimal_quintic [skip-build]
Adds a max_iter parameter (default 1000) to replace the open-ended while loop. Raises ValueError if the linear solve or node update produces non-finite values, and RuntimeError if convergence is not reached within max_iter iterations.206b280
-
Apply pre-commit formatting (yapf, isort) [skip-build]
aff01db -
Add comment explaining _coeffs_list and Polar Express vs former NS [skip-build]
Documents what _coeffs_list is (precomputed Polar Express coefficients, minimax-optimal via Remez/equioscillation), and contrasts with the former hardcoded NS coefficients: former produced US'V^T with scattered singular values; Polar Express converges to the exact polar factor UV^T. Also removes unused loguru import. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>abaa449
-
Replace hardcoded NS coefficients with analytically optimal ones [skip-build]
Replaces the five fixed (a, b, c) quintic coefficients in `_zeropower_via_newtonschulz5` with per-step optimal coefficients computed by `_optimal_composition` / `_optimal_quintic` via the equioscillation (Chebyshev) criterion. The precomputed `_coeffs_list` (10 iters, l=1e-3, safety_factor_eps=1e-2, cushion=0.02) minimizes the worst-case deviation from 1 over the singular-value interval, and the function now supports arbitrary `steps` by repeating the last coefficient set. Docstrings updated accordingly. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>573242f