fix: resolve rope_theta from wherever the config stores it - #166
Merged
Conversation
Every Eagle3 draft config in `configs/draft_models/` builds its rotary embedding at base 10000 regardless of what it declares: `kimi_k25_eagle3` and `kimi_k25_eagle3_mla` ask for 50000, `qwen3_8b_eagle3` and `qwen3_8b_eagle3_mla` for 1000000, `minimax_m25_eagle3` for 5000000. Every one of them silently trains against frequencies for 10000, so each draft learns positional structure its target does not have. The cause is not in these configs. transformers 5.x moved `rope_theta` into `rope_parameters`, and both `_init_rope` implementations read it as a top-level attribute with `getattr(self.config, "rope_theta", 10000)`. For any config the library manages the attribute no longer exists, so the default wins — a rename downstream turning into wrong numerics here, with nothing raised and nothing logged. `K3DSparkConfig` is the sole config class that escapes, because it lifts the nested value back onto the attribute itself. `tests/test_eagle3_loss.py::TestRotaryConfigWiring::test_yarn_uses_rope_theta_as_base` already asserted `rotary.base == 50000.0` and has been failing, so the intended behaviour was never in doubt; the failure was simply invisible among the pre-existing environment failures in this suite. It passes again with this commit. Add `resolve_rope_theta` beside the existing rope normalization in `config/utils.py` and use it from all three rotary construction paths — `LlamaAttention`, `DeepSeekMLAAttention` and `DFlashAttention`. The top-level attribute is preferred and the `rope_parameters` entry is the fallback, which keeps this repo's own configs (`DFlashConfig` and friends, which declare the attribute and carry no `rope_parameters`) resolving exactly as they do today. No config in the tree declares both, and if one ever does with two different values the helper raises rather than picking a winner, since a stale legacy value sitting beside an updated block is precisely how a five-fold frequency error goes unnoticed. `generate_draft_model_config` needs the same treatment for the same reason: it copies `rope_theta` from the target via `hasattr`, which under transformers 5.x silently copies nothing, so every auto-generated draft config inherited the default too. The sweep in `tests/test_draft_rope.py` is the part that survives the next rename: it asserts every config under `configs/draft_models/` resolves to the base its file declares, so a third location for this field fails the suite instead of quietly degrading training. Alongside it, one test per rotary path pins the constructed base to the configured value. Verified against the whole suite in the patched vLLM image: failures are identical to the branch point, 115 before and after, all pre-existing (missing sglang, `torch.compile` rejecting this host's `-march`, Ray teardown), with the one previously-failing rotary test now passing. All six shipped draft configs now build at their declared base. Drafts trained before this commit used base 10000 and are not comparable to ones trained after it. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ec5d9ddf51
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| params = getattr(config, "rope_parameters", None) | ||
| nested = params.get("rope_theta") if isinstance(params, dict) else None | ||
|
|
||
| if top_level is not None and nested is not None and float(top_level) != float(nested): |
There was a problem hiding this comment.
When a DFlash/DSpark draft config is written in the new rope_parameters form without a top-level rope_theta, DFlashConfig still creates self.rope_theta = 10000.0 from its constructor default before preserving the nested field from **kwargs. This new conflict branch then raises even though the file only declared one value, so loading such a v5-style draft config fails instead of resolving the nested theta. Consider ignoring the class default when the top-level value was not explicitly provided, or normalizing these config classes before calling this helper.
Useful? React with 👍 / 👎.
The omission fixed two commits back was structural rather than careless. Each attention block carried its own copy of the same six-branch ladder over `rope_scaling`, and the copies had already drifted: the MLA one dropped `base` on the YaRN branch, skipped the `factor` validation its sibling performs for linear and dynamic scaling, had no `mrope` branch, and kept a dead `scaling_type in (None, "default")` arm the enclosing condition already excluded. Nothing detects that kind of divergence, because each copy is locally plausible. Collapse both into `build_rotary_embedding(config, dim, max_position_embeddings)` next to the rotary classes it constructs. `dim` is the only thing that ever differed between the two: the head dim for the GQA block, the rope side dim for the MLA one. `resolve_rope_theta` is called inside the builder, so a caller cannot forget to pass the base — the specific mistake this replaces. Both `_init_rope` methods stay as one-line wrappers, since subclasses call them. Reconciling the two copies takes the stricter behaviour in each case: linear and dynamic scaling now raise on a missing `factor` for MLA drafts too, rather than passing `None` into the embedding and failing later somewhere less obvious, and MLA configs gain the `mrope` branch they never had. Neither changes any config in this tree, which uses yarn or no scaling at all. `rope_config_get` replaces the two identical accessors — a module-level function in `deepseek_eagle` and a closure in `llama3_eagle` — that both existed only to read a key from a dict-or-object `rope_scaling`. Net effect is 145 lines deleted for 87 added, and the file-level TODO noting that `_init_rope` was "near-identical to LlamaAttention._init_rope (~60 lines)" is resolved rather than restated. Verified in the patched vLLM image: failures are identical before and after, 115 either way, with no test moving in either direction; `test_draft_rope.py` and `test_dspark.py` pass, as does the previously-failing `TestRotaryConfigWiring::test_yarn_uses_rope_theta_as_base`; `ruff check` and `ruff format --check` clean. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
torchspec-bot
force-pushed
the
export/vllm-pp-hidden-states
branch
from
August 9, 2026 21:55
fb4eb7f to
c8314a5
Compare
torchspec-bot
force-pushed
the
export/rope-theta
branch
from
August 9, 2026 21:55
ec5d9dd to
7adce24
Compare
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.
Every Eagle3 draft config in
configs/draft_models/builds its rotary embedding at base 10000 no matter what it declares:kimi_k25_eagle3kimi_k25_eagle3_mlaqwen3_8b_eagle3qwen3_8b_eagle3_mlaminimax_m25_eagle3kimi_k3_dspark_mlatransformers 5.x moved
rope_thetaintorope_parameters, and both_init_ropeimplementations read it as a top-level attribute withgetattr(self.config, "rope_theta", 10000). For any config the library manages that attribute no longer exists, so the default wins silently.K3DSparkConfigis the only config class that escapes, because it lifts the nested value back onto the attribute itself.tests/test_eagle3_loss.py::TestRotaryConfigWiring::test_yarn_uses_rope_theta_as_basealready assertedrotary.base == 50000.0and has been failing — the intent was never ambiguous, the failure was just invisible among this suite's pre-existing environment failures. It passes again here.What changed
resolve_rope_thetainconfig/utils.py, beside the existing rope normalization, used by all three rotary paths (LlamaAttention,DeepSeekMLAAttention,DFlashAttention). Top-level attribute preferred,rope_parametersas fallback, so this repo's ownDFlashConfig-style configs resolve exactly as they do today. No config declares both; if one ever does with differing values it raises rather than picking a winner.generate_draft_model_configcopiedrope_thetafrom the target viahasattr, which under 5.x copied nothing, so auto-generated draft configs inherited the default too.tests/test_draft_rope.pysweeps every config underconfigs/draft_models/and asserts it resolves to the base its file declares, so a third location for this field fails the suite instead of quietly degrading training. Plus one test per rotary path.Why it could happen at all
Each attention block carried its own copy of the same six-branch ladder over
rope_scaling, and the copies had already drifted — the MLA one droppedbaseon the YaRN branch, skipped thefactorvalidation its sibling does for linear and dynamic scaling, had nomropebranch, and kept a dead arm the enclosing condition already excluded. The second commit collapses both intobuild_rotary_embedding(config, dim, max_position_embeddings), wheredimis the only thing that ever differed between callers.resolve_rope_thetais called inside the builder, so a caller cannot forget to pass the base — the exact mistake being replaced. Net 145 lines deleted for 87 added, and the file-level TODO about_init_ropebeing "near-identical to LlamaAttention._init_rope (~60 lines)" is resolved rather than restated.Reconciling the copies takes the stricter behaviour each way: MLA drafts now raise on a missing
factorinstead of passingNonedown, and gain themropebranch. No config in this tree is affected — they all use yarn or no scaling.Base branch
Stacked on
export/vllm-pp-hidden-states(#165) rather thanmain. The MLA half of this fix only takes effect together with the one-linebase=rope_thetain that PR'sfix(models): build the YaRN rotary cache at the configured rope_theta, and duplicating that line here would conflict on merge. Happy to retarget atmainonce #165 lands.Test plan
tests/test_draft_rope.py: 8 tests, 6 subtests (one per shipped config) passtorch.compilerejecting this host's-march, Ray teardown)test_yarn_uses_rope_theta_as_basenow passesruff checkandruff format --checkclean across 143 filesConsequence
Drafts trained before this used base 10000 and are not comparable to ones trained after it.