feat: recipe arms — headwise Muon + Sinkhorn embedding update - #225
Merged
Merged
Conversation
Two single-variable arms off the 2.1 recipe (run-007, 4.5701), DeepSeek-V4.1-Flash sec 2.5 / Alg 1: - ara-diac-small-2-1-hwmuon: head-wise Muon on Q/K (existing flag, first recipe run) -> run-014-hwmuon - ara-diac-small-2-1-skembed: sinkhorn_embed routes the 2D embedding-like tensors (tied table, prediction heads) from AdamW to SinkhornUpdate via the existing side-opt hook -> run-015-skembed muon.embed_named extracts the 2D embedding-like predicate so the selector and split_parameters share one classification (specs cover the MECE contract: exactly split's 2D AdamW set, never Muon's). Also fixes an engram-side routing bug: split_parameters handed _engram.table.weight to Muon while the side-opt stepped it too — double-stepped in run-013-engram. Tables are now excluded from muon_params (that run's FLAT verdict stands as measured; this corrects future engram arms).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Two single-variable recipe arms off 2.1 (run-007, 4.5701) — the DeepSeek-V4.1-Flash levers measured on our stack (TODO.impl/03 and /06's "landed awaiting run" items):
ara-diac-small-2-1-hwmuon): per-slice Newton–Schulz on Q/K projections (headwise_muonflag existed; this is its first recipe run) →run-014-hwmuonara-diac-small-2-1-skembed): newsinkhorn_embedflag routes the 2D embedding-like tensors (tied table, prediction heads) from AdamW toSinkhornUpdatevia the existing_engram_table_optside-opt hook →run-015-skembedCode
muon.embed_named: the 2D embedding-like selector, extracting the predicatesplit_parametersalready used so both share one classification. Specs assert the MECE contract: exactly split's 2D AdamW set, never Muon's; 1D layer norms excluded.split_parameterswas handing_engram.table.weightto Muon while the Sinkhorn side-opt also stepped it — double-stepped in the completedrun-013-engram(its FLAT verdict stands as measured with that wiring; this corrects future engram arms).Run protocol
Canonical r7 labels (sha e70ce991) pre-seeded into both out dirs;
labels_completetrusted. Launch:modal run --detach src/gpu/modal_distill.py::main --spec ara-diac-small-2-1-hwmuon --epochs 6(and...-skembed), gate each vs 4.5701 full-set at eval.