From d3b8342c1e9da3f02acb6fdff98ad8df1d9e5b00 Mon Sep 17 00:00:00 2001 From: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:39:53 +0000 Subject: [PATCH 1/2] fix(models): resolve rope_theta from wherever the config stores it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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> --- tests/test_draft_rope.py | 147 +++++++++++++++++++++++ torchspec/config/utils.py | 42 +++++++ torchspec/models/draft/deepseek_eagle.py | 3 +- torchspec/models/draft/dflash.py | 4 +- torchspec/models/draft/llama3_eagle.py | 8 +- 5 files changed, 199 insertions(+), 5 deletions(-) create mode 100644 tests/test_draft_rope.py diff --git a/tests/test_draft_rope.py b/tests/test_draft_rope.py new file mode 100644 index 00000000..df839b02 --- /dev/null +++ b/tests/test_draft_rope.py @@ -0,0 +1,147 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Tests that a draft's rotary base is the one its config declares. + +transformers 5.x moved `rope_theta` from a top-level config attribute into +`rope_parameters`, which silently reduced every affected draft to the library +default of 10000 while its config advertised something else. The sweep over +`configs/draft_models/` is the part that fails if a future release moves the +field again. +""" + +import glob +import json +import os +import unittest + +from transformers import LlamaConfig +from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + +from torchspec.config.utils import resolve_rope_theta +from torchspec.models.draft.auto import AutoDraftModelConfig +from torchspec.models.draft.deepseek_eagle import DeepSeekMLAAttention +from torchspec.models.draft.dflash import DFlashAttention, DFlashConfig +from torchspec.models.draft.llama3_eagle import LlamaAttention + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DRAFT_CONFIG_DIR = os.path.join(REPO_ROOT, "configs", "draft_models") + +YARN = { + "rope_type": "yarn", + "factor": 32.0, + "original_max_position_embeddings": 64, + "beta_fast": 32, + "beta_slow": 1, + "mscale": 1.0, + "mscale_all_dim": 1.0, +} + + +def _declared_theta(raw): + """The base a config file asks for, from wherever the file happens to put it.""" + block = raw.get("rope_parameters") or raw.get("rope_scaling") or {} + nested = block.get("rope_theta") + return nested if nested is not None else raw.get("rope_theta") + + +class TestResolveRopeTheta(unittest.TestCase): + def test_prefers_top_level_attribute(self): + cfg = DFlashConfig(rope_theta=50000.0) + self.assertEqual(resolve_rope_theta(cfg), 50000.0) + + def test_falls_back_to_rope_parameters(self): + # transformers 5.x absorbs a top-level rope_theta into rope_parameters, + # leaving no attribute for the naive read to find. + cfg = LlamaConfig(rope_theta=1000000.0) + self.assertIsNone(getattr(cfg, "rope_theta", None)) + self.assertEqual(resolve_rope_theta(cfg), 1000000.0) + + def test_default_when_declared_nowhere(self): + cfg = DFlashConfig() + del cfg.rope_theta + self.assertEqual(resolve_rope_theta(cfg), 10000.0) + self.assertIsNone(resolve_rope_theta(cfg, default=None)) + + def test_conflicting_values_are_refused(self): + cfg = LlamaConfig(rope_theta=1000000.0) + cfg.rope_theta = 10000.0 # stale legacy value left beside the block + with self.assertRaisesRegex(ValueError, "Conflicting rope_theta"): + resolve_rope_theta(cfg) + + +class TestShippedDraftConfigs(unittest.TestCase): + """Every config in configs/draft_models/ must resolve to the base it declares.""" + + def test_declared_theta_is_what_resolves(self): + paths = sorted(glob.glob(os.path.join(DRAFT_CONFIG_DIR, "*.json"))) + self.assertGreater(len(paths), 0, "no draft configs found") + for path in paths: + with self.subTest(config=os.path.basename(path)): + declared = _declared_theta(json.load(open(path))) + if declared is None: + self.skipTest("config declares no rope_theta") + cfg = AutoDraftModelConfig.from_file(path) + self.assertEqual(resolve_rope_theta(cfg), float(declared)) + + +class TestAttentionUsesResolvedTheta(unittest.TestCase): + """Each rotary construction path must reach the resolved base, not the default.""" + + def test_llama_yarn_path(self): + cfg = LlamaConfig( + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=4, + max_position_embeddings=2048, + rope_theta=1000000.0, + rope_parameters=dict(YARN), + ) + cfg.target_hidden_size = 64 + self.assertEqual(LlamaAttention(cfg).rotary_emb.base, 1000000.0) + + def test_deepseek_mla_yarn_path(self): + cfg = DeepseekV3Config( + hidden_size=64, + num_attention_heads=4, + q_lora_rank=32, + kv_lora_rank=16, + qk_nope_head_dim=16, + qk_rope_head_dim=8, + v_head_dim=16, + max_position_embeddings=2048, + rope_theta=50000.0, + rope_parameters=dict(YARN), + ) + self.assertEqual(DeepSeekMLAAttention(cfg).rotary_emb.base, 50000.0) + + def test_dflash_path(self): + cfg = DFlashConfig( + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=512, + rope_theta=1000000.0, + ) + self.assertEqual(DFlashAttention(cfg).rotary_emb.base, 1000000.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchspec/config/utils.py b/torchspec/config/utils.py index 630cd4f9..6025a0b6 100644 --- a/torchspec/config/utils.py +++ b/torchspec/config/utils.py @@ -22,6 +22,7 @@ import json import logging import warnings +from typing import Optional import torch from transformers import AutoConfig, AutoTokenizer @@ -58,6 +59,39 @@ def _normalize_rope_scaling(rope_scaling): return normalized +def resolve_rope_theta(config, default: Optional[float] = 10000.0) -> Optional[float]: + """Return the RoPE base a config asks for, wherever it happens to be stored. + + transformers 5.x keeps `rope_theta` inside `rope_parameters` instead of as a + top-level attribute, so reading the attribute alone silently yields `default` + for any config loaded from such a checkpoint — a draft then trains against + frequencies for 10000 while its config declares something else entirely. + + Configs written by this repo (`DFlashConfig` and friends) declare the + attribute themselves and have no `rope_parameters`, so the attribute is + preferred and those paths are unaffected. The two disagree only if a stale + legacy value was left behind next to an updated block, which is ambiguous + enough to refuse rather than silently resolve. + + Pass `default=None` to distinguish "declared nowhere" from a real value. + """ + top_level = getattr(config, "rope_theta", None) + 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): + raise ValueError( + f"Conflicting rope_theta: attribute is {top_level} but rope_parameters " + f"carries {nested}. Remove the stale top-level value, or set both to the " + "base the checkpoint was trained with." + ) + if top_level is not None: + return float(top_level) + if nested is not None: + return float(nested) + return None if default is None else float(default) + + def generate_draft_model_config( target_model_path: str, template_config_path: str = None, cache_dir: str = None ): @@ -139,6 +173,14 @@ def generate_draft_model_config( value = _normalize_rope_scaling(value) draft_config[draft_param] = value + # The hasattr copy above cannot see a rope_theta that lives only inside the + # target's rope_parameters, which is where transformers 5.x puts it. + for source in (text_config, target_config): + theta = resolve_rope_theta(source, default=None) + if theta is not None: + draft_config["rope_theta"] = theta + break + draft_config["num_hidden_layers"] = 1 draft_config["tie_word_embeddings"] = False draft_config["use_cache"] = True diff --git a/torchspec/models/draft/deepseek_eagle.py b/torchspec/models/draft/deepseek_eagle.py index c38cd11a..11eed395 100644 --- a/torchspec/models/draft/deepseek_eagle.py +++ b/torchspec/models/draft/deepseek_eagle.py @@ -30,6 +30,7 @@ from torch.nn.attention.flex_attention import flex_attention from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from torchspec.config.utils import resolve_rope_theta from torchspec.models.draft.base import Eagle3DraftModel # TODO: Extract shared components into a common module to reduce duplication: @@ -160,7 +161,7 @@ def _init_rope(self): """Initialize rotary embeddings with qk_rope_head_dim as the dimension.""" rope_dim = self.qk_rope_head_dim rope_scaling = self.config.rope_scaling - rope_theta = getattr(self.config, "rope_theta", 10000) + rope_theta = resolve_rope_theta(self.config) rget = partial(_rope_config_get, rope_scaling) scaling_type = rget("rope_type", rget("type")) diff --git a/torchspec/models/draft/dflash.py b/torchspec/models/draft/dflash.py index 37041104..21a14e3c 100644 --- a/torchspec/models/draft/dflash.py +++ b/torchspec/models/draft/dflash.py @@ -38,6 +38,8 @@ from safetensors import safe_open from transformers import PretrainedConfig, PreTrainedModel +from torchspec.config.utils import resolve_rope_theta + class DFlashConfig(PretrainedConfig): """Configuration for DFlash draft model.""" @@ -189,7 +191,7 @@ def __init__(self, config: PretrainedConfig): self.rotary_emb = DFlashRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, - base=getattr(config, "rope_theta", 10000.0), + base=resolve_rope_theta(config), ) def forward( diff --git a/torchspec/models/draft/llama3_eagle.py b/torchspec/models/draft/llama3_eagle.py index 99436c7e..22dd2244 100644 --- a/torchspec/models/draft/llama3_eagle.py +++ b/torchspec/models/draft/llama3_eagle.py @@ -30,6 +30,7 @@ from transformers.activations import ACT2FN from transformers.models.llama.configuration_llama import LlamaConfig +from torchspec.config.utils import resolve_rope_theta from torchspec.models.draft.base import Eagle3DraftModel from torchspec.models.ops.flex_attention import ( compile_friendly_flex_attention, @@ -1156,6 +1157,7 @@ def __init__(self, config): def _init_rope(self): rope_scaling = self.config.rope_scaling + rope_theta = resolve_rope_theta(self.config) def rope_get(key, default=None): if rope_scaling is None: @@ -1170,7 +1172,7 @@ def rope_get(key, default=None): self.rotary_emb = LlamaRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, - base=getattr(self.config, "rope_theta", 10000), + base=rope_theta, ) else: scaling_factor = rope_get("factor") @@ -1200,7 +1202,7 @@ def rope_get(key, default=None): self.rotary_emb = LlamaRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, - base=getattr(self.config, "rope_theta", 10000), + base=rope_theta, scaling_factor=(scaling_factor if scaling_factor is not None else 1.0), low_freq_factor=rope_get("low_freq_factor"), high_freq_factor=rope_get("high_freq_factor"), @@ -1214,7 +1216,7 @@ def rope_get(key, default=None): self.rotary_emb = LlamaYarnRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, - base=getattr(self.config, "rope_theta", 10000), + base=rope_theta, original_max_position_embeddings=rope_get("original_max_position_embeddings"), scaling_factor=scaling_factor, beta_fast=rope_get("beta_fast"), From 7adce248dff9d2b089734c6d7765e301c5450bab Mon Sep 17 00:00:00 2001 From: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:55:08 +0000 Subject: [PATCH 2/2] refactor(models): build every draft's rotary embedding in one place MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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/models/draft/deepseek_eagle.py | 82 ++----------- torchspec/models/draft/llama3_eagle.py | 150 ++++++++++++----------- 2 files changed, 87 insertions(+), 145 deletions(-) diff --git a/torchspec/models/draft/deepseek_eagle.py b/torchspec/models/draft/deepseek_eagle.py index 11eed395..3619239b 100644 --- a/torchspec/models/draft/deepseek_eagle.py +++ b/torchspec/models/draft/deepseek_eagle.py @@ -30,24 +30,20 @@ from torch.nn.attention.flex_attention import flex_attention from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config -from torchspec.config.utils import resolve_rope_theta from torchspec.models.draft.base import Eagle3DraftModel # TODO: Extract shared components into a common module to reduce duplication: # - LlamaMLP, LlamaRMSNorm, RoPE classes → torchspec/models/draft/modules.py -# - _init_rope() is near-identical to LlamaAttention._init_rope (~60 lines) # - DecoderLayer.forward() is line-for-line identical to LlamaDecoderLayer.forward() # - embed_input_ids/project_hidden_states/compute_logits/backbone are identical # to LlamaForCausalLMEagle3 and could live in Eagle3DraftModel base class # - Suffix attention loop could be batched (einsum instead of Python loop) for # both this file and llama3_eagle.py from torchspec.models.draft.llama3_eagle import ( - LlamaDynamicNTKScalingRotaryEmbedding, - LlamaLinearScalingRotaryEmbedding, LlamaMLP, LlamaRMSNorm, - LlamaRotaryEmbedding, - LlamaYarnRotaryEmbedding, + build_rotary_embedding, + rope_config_get, yarn_get_mscale, ) from torchspec.models.ops.flex_attention import ( @@ -56,16 +52,6 @@ ) from torchspec.utils.logging import logger, print_with_rank - -def _rope_config_get(rope_scaling, key, default=None): - """Get a value from rope_scaling config (dict or object).""" - if rope_scaling is None: - return default - if isinstance(rope_scaling, dict): - return rope_scaling.get(key, default) - return getattr(rope_scaling, key, default) - - # ── Interleaved RoPE (DeepSeek convention) ──────────────────────────────── # # DeepSeek MLA uses interleaved-pair rotation where consecutive dimension @@ -158,70 +144,16 @@ def __init__(self, config: DeepseekV3Config): self.softmax_scale = self._compute_softmax_scale() def _init_rope(self): - """Initialize rotary embeddings with qk_rope_head_dim as the dimension.""" - rope_dim = self.qk_rope_head_dim - rope_scaling = self.config.rope_scaling - rope_theta = resolve_rope_theta(self.config) - rget = partial(_rope_config_get, rope_scaling) - scaling_type = rget("rope_type", rget("type")) - - if rope_scaling is None or scaling_type == "default": - self.rotary_emb = LlamaRotaryEmbedding( - rope_dim, - max_position_embeddings=self.max_position_embeddings, - base=rope_theta, - ) - else: - scaling_factor = rget("factor") - - if scaling_type in (None, "default"): - self.rotary_emb = LlamaRotaryEmbedding( - rope_dim, - max_position_embeddings=self.max_position_embeddings, - base=rope_theta, - ) - elif scaling_type == "linear": - self.rotary_emb = LlamaLinearScalingRotaryEmbedding( - rope_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - ) - elif scaling_type == "dynamic": - self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding( - rope_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - ) - elif scaling_type == "llama3": - self.rotary_emb = LlamaRotaryEmbedding( - rope_dim, - max_position_embeddings=self.max_position_embeddings, - base=rope_theta, - scaling_factor=scaling_factor if scaling_factor is not None else 1.0, - low_freq_factor=rget("low_freq_factor"), - high_freq_factor=rget("high_freq_factor"), - orig_max_position=rget("original_max_position_embeddings"), - ) - elif scaling_type == "yarn": - self.rotary_emb = LlamaYarnRotaryEmbedding( - rope_dim, - max_position_embeddings=self.max_position_embeddings, - base=rope_theta, - original_max_position_embeddings=rget("original_max_position_embeddings"), - scaling_factor=scaling_factor, - beta_fast=rget("beta_fast"), - beta_slow=rget("beta_slow"), - mscale=rget("mscale"), - mscale_all_dim=rget("mscale_all_dim"), - ) - else: - raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + """Rotate only the rope side dims, so the cache is built at qk_rope_head_dim.""" + self.rotary_emb = build_rotary_embedding( + self.config, self.qk_rope_head_dim, self.max_position_embeddings + ) def _compute_softmax_scale(self) -> float: """Compute softmax scale, incorporating YaRN mscale if applicable.""" rope_scaling = self.config.rope_scaling if rope_scaling is not None: - rget = partial(_rope_config_get, rope_scaling) + rget = partial(rope_config_get, rope_scaling) scaling_type = rget("rope_type", rget("type")) if scaling_type == "yarn": factor = rget("factor", 1.0) diff --git a/torchspec/models/draft/llama3_eagle.py b/torchspec/models/draft/llama3_eagle.py index 22dd2244..33e70aad 100644 --- a/torchspec/models/draft/llama3_eagle.py +++ b/torchspec/models/draft/llama3_eagle.py @@ -20,6 +20,7 @@ import math import os +from functools import partial from typing import Optional, Tuple import torch @@ -605,6 +606,82 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): ) +def rope_config_get(rope_scaling, key, default=None): + """Read a key out of a rope_scaling config, which may be a dict or an object.""" + if rope_scaling is None: + return default + if isinstance(rope_scaling, dict): + return rope_scaling.get(key, default) + return getattr(rope_scaling, key, default) + + +def build_rotary_embedding(config, dim: int, max_position_embeddings: int): + """Build the rotary embedding a config asks for. + + Every draft that uses these rotary classes goes through here, and `dim` is the + only thing that differs between callers — the head dim for GQA blocks, the + rope side dim for MLA ones. Each attention block used to carry its own copy of + this branch ladder, and the copies drifted: the MLA one dropped `base` on the + YaRN branch, so drafts trained at 10000 while their config declared otherwise. + Adding a scaling type or an argument in one place now reaches every draft. + """ + rope_scaling = config.rope_scaling + rope_theta = resolve_rope_theta(config) + rget = partial(rope_config_get, rope_scaling) + scaling_type = rget("rope_type", rget("type")) + + if rope_scaling is None or scaling_type in (None, "default"): + return LlamaRotaryEmbedding( + dim, + max_position_embeddings=max_position_embeddings, + base=rope_theta, + ) + + scaling_factor = rget("factor") + + if scaling_type == "linear": + if scaling_factor is None: + raise ValueError("Linear RoPE scaling requires 'factor' in rope_scaling config.") + return LlamaLinearScalingRotaryEmbedding( + dim, + max_position_embeddings=max_position_embeddings, + scaling_factor=scaling_factor, + ) + if scaling_type == "dynamic": + if scaling_factor is None: + raise ValueError("Dynamic RoPE scaling requires 'factor' in rope_scaling config.") + return LlamaDynamicNTKScalingRotaryEmbedding( + dim, + max_position_embeddings=max_position_embeddings, + scaling_factor=scaling_factor, + ) + if scaling_type == "llama3": + return LlamaRotaryEmbedding( + dim, + max_position_embeddings=max_position_embeddings, + base=rope_theta, + scaling_factor=scaling_factor if scaling_factor is not None else 1.0, + low_freq_factor=rget("low_freq_factor"), + high_freq_factor=rget("high_freq_factor"), + orig_max_position=rget("original_max_position_embeddings"), + ) + if scaling_type == "mrope": + return LlamaMutiRotaryEmbedding(dim, max_position_embeddings=max_position_embeddings) + if scaling_type == "yarn": + return LlamaYarnRotaryEmbedding( + dim, + max_position_embeddings=max_position_embeddings, + base=rope_theta, + original_max_position_embeddings=rget("original_max_position_embeddings"), + scaling_factor=scaling_factor, + beta_fast=rget("beta_fast"), + beta_slow=rget("beta_slow"), + mscale=rget("mscale"), + mscale_all_dim=rget("mscale_all_dim"), + ) + raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + + _SNAP_Q = 128 # Alignment granularity for Q_LEN bucketing. @@ -1156,76 +1233,9 @@ def __init__(self, config): self._init_rope() def _init_rope(self): - rope_scaling = self.config.rope_scaling - rope_theta = resolve_rope_theta(self.config) - - def rope_get(key, default=None): - if rope_scaling is None: - return default - if isinstance(rope_scaling, dict): - return rope_scaling.get(key, default) - return getattr(rope_scaling, key, default) - - scaling_type = rope_get("rope_type", rope_get("type")) - - if rope_scaling is None or scaling_type == "default": - self.rotary_emb = LlamaRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - base=rope_theta, - ) - else: - scaling_factor = rope_get("factor") - - if scaling_type == "linear": - if scaling_factor is None: - raise ValueError( - "Linear RoPE scaling requires 'factor' in rope_scaling config." - ) - self.rotary_emb = LlamaLinearScalingRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - ) - elif scaling_type == "dynamic": - if scaling_factor is None: - raise ValueError( - "Dynamic RoPE scaling requires 'factor' in rope_scaling config." - ) - self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - ) - elif scaling_type == "llama3": - # for nv type - self.rotary_emb = LlamaRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - base=rope_theta, - scaling_factor=(scaling_factor if scaling_factor is not None else 1.0), - low_freq_factor=rope_get("low_freq_factor"), - high_freq_factor=rope_get("high_freq_factor"), - orig_max_position=rope_get("original_max_position_embeddings"), - ) - elif scaling_type == "mrope": - self.rotary_emb = LlamaMutiRotaryEmbedding( - self.head_dim, max_position_embeddings=self.max_position_embeddings - ) - elif scaling_type == "yarn": - self.rotary_emb = LlamaYarnRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - base=rope_theta, - original_max_position_embeddings=rope_get("original_max_position_embeddings"), - scaling_factor=scaling_factor, - beta_fast=rope_get("beta_fast"), - beta_slow=rope_get("beta_slow"), - mscale=rope_get("mscale"), - mscale_all_dim=rope_get("mscale_all_dim"), - ) - else: - raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + self.rotary_emb = build_rotary_embedding( + self.config, self.head_dim, self.max_position_embeddings + ) def forward( self,