fix(dspark): initialize and optimize the Markov lookup embedding as an embedding - #971
fix(dspark): initialize and optimize the Markov lookup embedding as an embedding#971fynnsu wants to merge 1 commit into
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (1)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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. Comment |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
|
The quality checks have failed. Please run |
b70c09e to
4819ff4
Compare
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
The change is well-motivated: markov_w1 is an embedding lookup table and should be treated as one for both initialization and optimizer routing. Initializing with std=0.01 keeps the Markov sequential bias small at the start of training (the default N(0, 1) would let W1 @ W2 dominate the DFlash logits from step 1), and routing to AdamW follows the standard Muon convention for embedding layers.
The commit message states "Add tests for both: the Muon/AdamW split had no coverage," but the diff contains no test files. Are those tests in a separate PR in the stack, or were they inadvertently left out?
🤖 Generated with Claude Code using the /pr-review skill
4819ff4 to
9759290
Compare
9759290 to
db3570e
Compare
There was a problem hiding this comment.
Theoretically this should be applied to both markov_w1 and markov_w2.
Please see my previous ablation study on why this is difficult.
#866.
We need more ablation study for this bug fix.
…n embedding MarkovHead.markov_w1 indexes the verifier vocabulary, but it is stored as a 2D nn.Embedding weight, so two pieces of code treated it as an ordinary matrix. split_named_params_for_muon routed it to Muon's orthogonalized update. Muon excludes embeddings and the output head by convention, since orthogonalizing a lookup table updates every row of the vocabulary on a step where only a few tokens appeared. Add markov_w1 to the AdamW name hints so it is grouped with embed_tokens and lm_head. It also kept nn.Embedding's default N(0, 1) initialization. markov_w2 projects the lookup straight onto the draft logits, so a unit-scale init lets the sequential bias dominate the DFlash logits from the first step. Initialize it with std=0.01. Add tests for both: the Muon/AdamW split had no coverage. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
db3570e to
b6d6735
Compare
3 participants
Purpose
Change the way we init the markov head and switch it to using the Adam optimizer. These should be better defaults for the embedding layer.
Tests
Not extensively validated, and change is primarily theoretical. More experimentation to follow.
Checklist
I have filled in:
Stack created with GitHub Stacks CLI • Give Feedback 💬