[baseten] fix(dsa): plumb indexer RoPE interleave convention from config - #6
Merged
Merged
Conversation
GLM-5.2 declares indexer_rope_interleave: true (GPT-J interleaved rope in the DSA indexer; vLLM reads it as is_neox_style = not flag). The indexer's _apply_rope hardcoded the DeepSeek-V3.2 convention (non-interleaved), so the trainer scored its indexer on differently-rotated q/k than serving engines. Top-k selections still overlap at small candidate pools but diverge progressively with sequence length: measured trainer<->vLLM per-token logprob KL grows 0.010 -> 0.084 between 9k and 15k tokens, making long-context RL unusable. With the convention plumbed (dsa_indexer_rope_interleave, set by the GLM bridge from the checkpoint config), the same probe measures 0.009 at 15k — flat in length, zero systematic bias. Default False preserves DeepSeek-V3.2 behavior exactly.
pstefa1707
force-pushed
the
fix/dsa-indexer-rope-interleave
branch
from
July 7, 2026 20:26
c5b095b to
377e4f5
Compare
pstefa1707
pushed a commit
to basetenlabs/Megatron-Bridge
that referenced
this pull request
Jul 7, 2026
…nfig GLM-5.2 declares indexer_rope_interleave: true — its DSA indexer uses GPT-J interleaved RoPE, which serving engines honor (vLLM: is_neox_style = not flag). The bridge never plumbed it, so Megatron's indexer defaulted to the DeepSeek-V3.2 non-interleaved convention and its top-k selection drifted from inference progressively with sequence length (trainer<->vLLM logprob KL 0.010 -> 0.084 between 9k and 15k tokens; 0.009 flat after the fix). Bumps 3rdparty/Megatron-LM to pick up dsa_indexer_rope_interleave (basetenlabs/Megatron-LM#6).
JackRao123
approved these changes
Jul 7, 2026
pstefa1707
pushed a commit
to basetenlabs/Megatron-Bridge
that referenced
this pull request
Jul 7, 2026
…nfig (#10) GLM-5.2 declares indexer_rope_interleave: true — its DSA indexer uses GPT-J interleaved RoPE, which serving engines honor (vLLM: is_neox_style = not flag). The bridge never plumbed it, so Megatron's indexer defaulted to the DeepSeek-V3.2 non-interleaved convention and its top-k selection drifted from inference progressively with sequence length (trainer<->vLLM logprob KL 0.010 -> 0.084 between 9k and 15k tokens; 0.009 flat after the fix). Bumps 3rdparty/Megatron-LM to pick up dsa_indexer_rope_interleave (basetenlabs/Megatron-LM#6).
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
2 participants
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.
Bug
GLM-5.2 sets
indexer_rope_interleave: truein its HF config — its DSA indexer applies GPT-J interleaved RoPE (vLLM honors this:is_neox_style = not config.indexer_rope_interleave). Megatron'sDSAIndexer._apply_ropehardcodedmla_rotary_interleaved=False(the DeepSeek-V3.2 convention), and the flag was plumbed nowhere.Result: the trainer scores its indexer on differently-rotated q/k than serving engines. Top-k selections still mostly agree while the candidate pool is small, then diverge progressively with sequence length — wrong attended-token sets compounding through all sparse layers.
Measured (GLM-5.2, 3×8 B200 trainer vs vLLM 0.23, teacher-forced per-token logprob KL on identical sequences):
After the fix the curve is flat in length and the systematic trainer-below-sampler bias (−0.09 nats/token at 15k) vanishes (−0.00002). Long-context GLM-5.2 RL goes from unusable (8/8 steps over prime-rl's 0.015 mismatch gate) to passing with margin (0/8).
Change
TransformerConfig.dsa_indexer_rope_interleave: bool = False— declared on the config (same conversion-survival rationale asdsa_indexer_topk_freq).DSAIndexer._apply_ropepasses it asmla_rotary_interleavedinstead of the hardcode.Falsepreserves DeepSeek-V3.2 behavior bit-for-bit.Companion Megatron-Bridge PR sets it from the checkpoint config in the GLM-5.2 bridge (
provider.dsa_indexer_rope_interleave = hf_config.indexer_rope_interleave).Diagnosis trail (three-way decode/prefill/trainer probes, length bisects, fp8-indexer counter-experiment) in Baseten TRN-1406.
🤖 Generated with Claude Code