fix(muon x dspark): Markov embedding should not use Muon

[WindChimeRan]

Purpose

Two defects in how DSpark's Markov head is initialized and optimized. They have to land together — fixing either alone leaves training worse than it is today, for the reason in the last section.

1. The head misses the weight-init sweep

DSparkDraftModel.__init__ attaches the Markov and confidence heads after super().__init__(), which already ran post_init() at the end of DFlash's constructor (dflash/core.py:124). The heads are never visited and keep raw torch defaults.

Measured on a real model:

parameter	as built	rest of the draft
markov_head.markov_w1.weight	1.0000	0.0200
markov_head.markov_w2.weight	0.0361	0.0200
layers.0.self_attn.q_proj.weight	0.0200	—

markov_w1 is an nn.Embedding, so it gets N(0, 1) — 50× the initializer_range every other draft matrix gets. The head's entire output is added to the draft logits (dspark/core.py:170) with nothing in between, so training opens against a random per-token bigram prior:

head type	markov_bias std	absmax	after
vanilla	0.577	3.34	0.006 / 0.04
gated	0.294	1.70	0.003 / 0.02
rnn	0.177	1.01	0.003 / 0.02

Calling post_init() once more at the end of __init__ fixes it. Modules the first pass covered carry _is_hf_initialized and are skipped, and the deliberately NaN-filled lm_head / embed_tokens / verifier_lm_head sentinels survive — verified, nothing but the two heads changes.

DeepSpec builds its Markov head before post_init() (modeling/dspark/qwen3/modeling.py:252 then :268), so it gets initializer_range and does not have this bug.

2. Muon orthogonalizes the Markov embedding

The Muon/AdamW split promises in its docstring that embeddings go to AdamW, but tests the parameter name: _ADAMW_NAME_HINTS = ("embed_tokens", "lm_head"). markov_head.markov_w1 is 2D with both dims > 1 and matches no hint, so it goes to Muon.

Muon replaces the momentum-smoothed gradient with its nearest orthogonal matrix, equalizing the update across singular directions. That is meaningful for a hidden Linear, where both axes are learned feature spaces. An embedding's row axis is a token id, so the table is a set of independent vectors rather than a linear map — orthogonalizing couples every token's update to every other token's gradient and flattens a spectrum that mostly encodes token frequency. Hence the convention this code already cites at optimizers.py:24.

Matching on isinstance(module, nn.Embedding) rather than adding markov_w1 to the tuple: name-matching is what created the gap, and extending it would re-arm the trap for the next embedding.

Why these cannot be split

The Muon group's effective decay is lr * wd = 1e-3 * 0.1 = 1e-4 per step (muon_lr defaults to 10 * lr). The AdamW group's is 1e-4 * 0.01 = 1e-6. 100×.

Most rows of a 150k-entry embedding never receive gradient on a given dataset, so decay is the only force acting on them. An untrained row starting at std 1.0:

steps	on main (Muon)	fix 2 alone (AdamW)	both fixes
10,000	0.368	0.990	0.020
30,000	0.050	0.970	0.020
100,000	0.000	0.905	0.018

On main the aggressive Muon decay drags untrained rows down to ~0.05 over a realistic run — accidentally near the 0.02 convention, masking the init bug. Correcting the routing alone removes that compensation while leaving the 50× init, so those rows stay at ~0.97 for the entire run and inject the full ±3.3-logit bias whenever such a token precedes. Reaching 0.05 under AdamW decay would take ~3M steps.

Tests

Two, one per defect, each red before its own change and green after.

test_markov_head_is_covered_by_the_weight_init_sweep — the Markov embedding's std must match a matrix the sweep certainly covered, rather than a hardcoded constant:

E  assert 0.9997901916503906 == 0.020039767026901245 ± 0.00601193

test_embeddings_route_to_adamw_regardless_of_name — an embedding under a name no hint matches must land in AdamW, while an ordinary matrix still reaches Muon:

E  AssertionError: assert 'markov_head.markov_w1.weight' in set()

pytest tests/unit -q
611 passed, 5 skipped

ruff check and ruff format --check pass over src/ and tests/.

Notes

Only fresh training is affected; any run loading a checkpoint overwrites these weights.

Two things I deliberately did not change:

• markov_head.markov_w2.weight is nn.Linear(markov_rank, draft_vocab_size), whose rows are indexed by the next token id and whose output is added to the logits — the output-side half of the same factorized token×token table, and arguably subject to the same exclusion as lm_head. Routing it would need a mechanism a model can use to declare "this matrix is token-indexed", since nothing about its type distinguishes it from q_proj. That deserves its own design rather than another name hint.
• Whether initializer_range is the right scale for a logit-bias head at all, versus starting it at zero so the head is a no-op at init. Matching DeepSpec is the conservative choice; near-zero is defensible.

Checklist

I have filled in:

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• (Optional) The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 5b4fa3c9-257a-4a5e-8ae8-82a380e5b371

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz