feat(gpu): headwise Muon, Sinkhorn-balanced update, DeepSeek measurement probes - #210
Merged
Merged
Conversation
added 2 commits
September 11, 2026 21:24
TODO.impl/03, /06 and the two measurement probes from the DeepSeek-V4.1-Flash learnings pass: - muon: headwise group (per-head Newton-Schulz for Q/K, config-gated by the headwise_muon spec flag, off by default), qk_named selector, shared momentum helper. Off-state routing bit-identical. - sinkhorn_update: Algorithm 1 as an optimizer for embedding/head matrices (Nesterov momentum, alternating row/col L2, near-zero row mask, sqrt(hidden) RMS conversion, gamma=0.18 lr correction, no weight decay, eps=1e-20 guards). - probe_speculative: draft/verify acceptance measurement over golden-v1 rows between the shipped layerdrop-int4 and small-2.1-int8 artifacts (early rows: acceptance 0.97-1.00, ~8.9 tokens/verify). - modal_teacher_sadeed: teacher-only SadeedDiac-25 predictions under the windowed protocol (used for the r6 per-domain slice, negative verdict recorded in rababa TODO.impl/02). - distill_specs: ara-diac-small-lite2 - the lite rung's untested variable, layer-drop init from the trained 2.1 student.
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.
The ml-models side of the DeepSeek-V4.1-Flash learnings pass (rababa TODO.impl/03, /06, and the two measurement probes).
Headwise Muon (
muon.py): per-head Newton-Schulz for Q/K projections - externally validated by DeepSeek-V4.1-Flash sec 2.5, GLM-5 and Kimi-K3. Newadd_headwise_group(converts an existing Muon group in place or adds one),qk_namedselector, shared_momentum_directionhelper. Config-gated by theheadwise_muonspec flag in the distill loop, off by default - off-state routing is bit-identical.Sinkhorn-balanced update (
sinkhorn_update.py): the report's Algorithm 1 as an optimizer for embedding/head matrices - Nesterov momentum, alternating row/column L2 normalization (odd K, ends rows), near-zero row masking at tau=1e-3, sqrt(hidden-dim) RMS conversion, gamma=0.18 lr correction, no weight decay. eps=1e-20 guards per the mHC log-domain lesson.Probes:
scripts/probe_speculative.py(draft/verify acceptance between the shipped layerdrop-int4 and small-2.1-int8 artifacts over golden-v1 rows; measured acceptance 0.97-1.00 at ~8.9 tokens per verifier pass, K=8; output-preserving greedy verification verified exact against the plain-path reference) andsrc/gpu/modal_teacher_sadeed.py(teacher-only SadeedDiac-25 predictions under the windowed protocol - produced the r6 per-domain slice behind the negative routing verdict).Spec:
ara-diac-small-lite2- the lite rung's one untested variable: layer-drop init from the trained 2.1 student instead of generic byt5-small, recipe otherwise identical to run-009 (launched separately).Tests:
test_muon_headwise.py(5) andtest_sinkhorn_update.py(6), all real tensors - per-slice equivalence, heads=1 identity, hand-computed balance case, manual algorithm match, no-weight-decay contract.