From ecda5f94ce0dfb723411d928df430765f82d908c Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Fri, 28 Aug 2026 17:55:31 +0800 Subject: [PATCH 1/4] Read rope_theta from rope_parameters in the Llama injection policy Kernel injection resolves rope_theta as if hasattr(self.policy.client_module.self_attn, 'config'): _config.rope_theta = ...self_attn.config.rope_theta else: _config.rope_theta = ...self_attn.rope_theta transformers 5.0 folded the rotary settings into config.rope_parameters and dropped the attribute, so against a stock LlamaConfig both branches raise. The first is taken -- LlamaAttention still has .config -- and injection dies with AttributeError: 'LlamaConfig' object has no attribute 'rope_theta' On transformers 5.8.0 that reproduces with a default LlamaConfig, so it is not specific to the DeepSeek checkpoint in #8340. Fall back to rope_parameters['rope_theta'] when the attribute is gone, keeping the old spellings first so pre-5.0 installs are unaffected. This is the same drift #7443 adapted to; the num_heads accessor it fixed alongside still resolves, so only this one moved. Fixes #8340 Signed-off-by: alanhuangyoo --- deepspeed/module_inject/containers/llama.py | 23 ++++++-- .../module_inject/test_llama_rope_theta.py | 57 +++++++++++++++++++ 2 files changed, 76 insertions(+), 4 deletions(-) create mode 100644 tests/unit/module_inject/test_llama_rope_theta.py diff --git a/deepspeed/module_inject/containers/llama.py b/deepspeed/module_inject/containers/llama.py index f57652aecdda..e9462da60164 100644 --- a/deepspeed/module_inject/containers/llama.py +++ b/deepspeed/module_inject/containers/llama.py @@ -20,6 +20,24 @@ ) +def _get_rope_theta(self_attn): + """Read rope_theta from whichever place the installed transformers keeps it. + + transformers < 5.0 exposes it as ``config.rope_theta``; 5.0 moved the rotary + settings into the ``rope_parameters`` dict and dropped the attribute, so the + older reads raise AttributeError against a stock LlamaConfig. Very old + versions kept it on the attention module itself. + """ + config = getattr(self_attn, 'config', None) + if config is not None: + if hasattr(config, 'rope_theta'): + return config.rope_theta + rope_parameters = getattr(config, 'rope_parameters', None) + if rope_parameters is not None and 'rope_theta' in rope_parameters: + return rope_parameters['rope_theta'] + return self_attn.rope_theta + + class DS_LLAMAContainer(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, BaseTransformerContainer): @@ -34,10 +52,7 @@ def create_module(self, config=None): _config.rotate_half = True _config.rotate_every_two = False _config.rotary_dim = self.hidden_size // self.num_attention_heads - if hasattr(self.policy.client_module.self_attn, 'config'): - _config.rope_theta = self.policy.client_module.self_attn.config.rope_theta - else: - _config.rope_theta = self.policy.client_module.self_attn.rope_theta + _config.rope_theta = _get_rope_theta(self.policy.client_module.self_attn) self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) return self.module diff --git a/tests/unit/module_inject/test_llama_rope_theta.py b/tests/unit/module_inject/test_llama_rope_theta.py new file mode 100644 index 000000000000..4e390e7076c5 --- /dev/null +++ b/tests/unit/module_inject/test_llama_rope_theta.py @@ -0,0 +1,57 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Llama kernel injection must find rope_theta wherever the installed transformers keeps it. + +The value has moved twice. Older releases exposed ``config.rope_theta``; transformers 5.0 +folded the rotary settings into ``config.rope_parameters`` and dropped the attribute, so an +injection policy that only knows the old spelling raises AttributeError against a stock +LlamaConfig. Older still, it lived on the attention module. +""" + +from types import SimpleNamespace + +import pytest + +from deepspeed.module_inject.containers.llama import _get_rope_theta + + +def test_reads_the_legacy_config_attribute(): + self_attn = SimpleNamespace(config=SimpleNamespace(rope_theta=500000.0)) + + assert _get_rope_theta(self_attn) == 500000.0 + + +def test_reads_rope_parameters_when_the_attribute_is_gone(): + # transformers >= 5.0: the attribute is absent and the value sits in the dict. + config = SimpleNamespace(rope_parameters={"rope_theta": 10000.0, "rope_type": "default"}) + self_attn = SimpleNamespace(config=config) + + assert _get_rope_theta(self_attn) == 10000.0 + + +def test_prefers_the_attribute_when_both_are_present(): + config = SimpleNamespace(rope_theta=500000.0, rope_parameters={"rope_theta": 10000.0}) + self_attn = SimpleNamespace(config=config) + + assert _get_rope_theta(self_attn) == 500000.0 + + +def test_falls_back_to_the_module_attribute(): + # No config at all, the layout the policy handled before configs were attached. + self_attn = SimpleNamespace(rope_theta=1000000.0) + + assert _get_rope_theta(self_attn) == 1000000.0 + + +def test_falls_back_when_rope_parameters_carries_no_theta(): + config = SimpleNamespace(rope_parameters={"rope_type": "default"}) + self_attn = SimpleNamespace(config=config, rope_theta=250000.0) + + assert _get_rope_theta(self_attn) == 250000.0 + + +def test_raises_when_nothing_carries_it(): + with pytest.raises(AttributeError): + _get_rope_theta(SimpleNamespace(config=SimpleNamespace())) From 3dd2716a181ab0f06d6ea6b8f55e34516d3f372a Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Fri, 4 Sep 2026 14:43:15 +0800 Subject: [PATCH 2/4] Pin the Llama rope_theta helper against a real LlamaConfig The six existing cases build the config out of SimpleNamespace, so nothing in the file pinned the claim the change rests on: which spelling a stock LlamaConfig actually carries. Every one of them passes against the previous implementation. Two added. The first resolves a non-default theta through the real class, so it stays meaningful on either side of the 5.0 boundary -- legacy attribute on 4.x, rope_parameters on 5.x -- and fails against the previous implementation, which raises AttributeError there on transformers 5.16.1: hasattr(LlamaConfig(rope_theta=5e5), 'rope_theta') -> False rope_parameters -> {'rope_theta': 500000.0, ...} The second guards the branch rather than the value. If transformers reinstates rope_theta as a deprecated property the first test still passes while the injection path silently changes branch, which only matters if the two spellings can disagree, so it asserts they do not. Signed-off-by: alanhuangyoo --- .../module_inject/test_llama_rope_theta.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/unit/module_inject/test_llama_rope_theta.py b/tests/unit/module_inject/test_llama_rope_theta.py index 4e390e7076c5..73dd2ee88af1 100644 --- a/tests/unit/module_inject/test_llama_rope_theta.py +++ b/tests/unit/module_inject/test_llama_rope_theta.py @@ -52,6 +52,40 @@ def test_falls_back_when_rope_parameters_carries_no_theta(): assert _get_rope_theta(self_attn) == 250000.0 +def test_resolves_against_a_real_llama_config(): + """The six cases above build the config by hand, so none of them pins the claim this + change rests on: which spelling a stock `LlamaConfig` actually carries. + + This one stays meaningful on either side of the 5.0 boundary — it takes the legacy + attribute on 4.x and `rope_parameters` on 5.x — and it uses a non-default theta, so a + helper that returned the class default would fail it. + """ + LlamaConfig = pytest.importorskip("transformers.models.llama.configuration_llama").LlamaConfig + + self_attn = SimpleNamespace(config=LlamaConfig(rope_theta=500000.0)) + + assert _get_rope_theta(self_attn) == 500000.0 + + +def test_the_two_spellings_do_not_disagree_on_a_real_config(): + """Guards the branch rather than the value. + + If transformers reinstates `rope_theta` as a deprecated property, the test above still + passes while the injection path silently changes which branch it takes. That is only a + problem if the two spellings can disagree, so this asserts they cannot. + """ + LlamaConfig = pytest.importorskip("transformers.models.llama.configuration_llama").LlamaConfig + config = LlamaConfig(rope_theta=500000.0) + + parameters = getattr(config, "rope_parameters", None) or {} + carried = [ + value for value in (getattr(config, "rope_theta", None), parameters.get("rope_theta")) if value is not None + ] + + assert carried, "neither spelling carries rope_theta on the installed transformers" + assert all(value == 500000.0 for value in carried), f"the spellings disagree: {carried}" + + def test_raises_when_nothing_carries_it(): with pytest.raises(AttributeError): _get_rope_theta(SimpleNamespace(config=SimpleNamespace())) From 8447987cd51a99f4cfb3334b5d39e92e9550bb7b Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 16:48:32 +0800 Subject: [PATCH 3/4] Refuse rotary variants kernel injection cannot serve @tohtana is right that this patch does not make DeepSeek-R1-Distill-Llama-8B work. It uses rope_type="llama3" with factor, low_freq_factor, high_freq_factor and original_max_position_embeddings, and only rope_theta reaches the kernel. The injected path has no place to put them. Its rotary embedding is built from a scalar base and nothing else: InferenceContext.get_rotary(rotary_dim, rope_theta) and DeepSpeedInferenceConfig carries rope_theta with no scaling fields at all. So after the crash fix that model would have started and run with unscaled positions, producing wrong output with no error. That is worse than the AttributeError this change exists to remove: the crash is at least visible. A config asking for a variant the kernel cannot implement is now refused, naming the type and pointing at running without kernel injection. Both spellings are covered, rope_parameters on 5.x and rope_scaling on 4.x. `default` and an absent type are unscaled and still resolve, so the configuration #8340 reported is fixed as before. Implementing llama3 scaled RoPE would mean propagating the parameters through DeepSpeedInferenceConfig and adding an inv_freq path and dispatch in the kernel. That is a feature rather than a crash fix and belongs in its own change. Signed-off-by: alanhuangyoo --- deepspeed/module_inject/containers/llama.py | 44 +++++++++++++--- .../module_inject/test_llama_rope_theta.py | 51 +++++++++++++++++++ 2 files changed, 89 insertions(+), 6 deletions(-) diff --git a/deepspeed/module_inject/containers/llama.py b/deepspeed/module_inject/containers/llama.py index e9462da60164..9dc2552c43cb 100644 --- a/deepspeed/module_inject/containers/llama.py +++ b/deepspeed/module_inject/containers/llama.py @@ -19,6 +19,23 @@ maybe_get_lora, ) +# The injected kernel builds its rotary embedding from a scalar base and nothing else +# (`InferenceContext.get_rotary(rotary_dim, rope_theta)`), so a config asking for a scaled +# variant cannot be honoured here. These are the spellings that mean "no scaling". +_UNSCALED_ROPE_TYPES = (None, 'default') + + +def _rope_type(config, rope_parameters): + """The rotary variant a config asks for, however the installed transformers spells it.""" + if rope_parameters: + rope_type = rope_parameters.get('rope_type', rope_parameters.get('type')) + if rope_type is not None: + return rope_type + scaling = getattr(config, 'rope_scaling', None) + if isinstance(scaling, dict): + return scaling.get('rope_type', scaling.get('type')) + return None + def _get_rope_theta(self_attn): """Read rope_theta from whichever place the installed transformers keeps it. @@ -27,14 +44,29 @@ def _get_rope_theta(self_attn): settings into the ``rope_parameters`` dict and dropped the attribute, so the older reads raise AttributeError against a stock LlamaConfig. Very old versions kept it on the attention module itself. + + A scaled variant is refused rather than silently reduced to its base. The kernel + implements a scalar theta only, so running one of these with just ``rope_theta`` + produces wrong positions with no error, which is worse than the AttributeError this + function exists to remove. """ config = getattr(self_attn, 'config', None) - if config is not None: - if hasattr(config, 'rope_theta'): - return config.rope_theta - rope_parameters = getattr(config, 'rope_parameters', None) - if rope_parameters is not None and 'rope_theta' in rope_parameters: - return rope_parameters['rope_theta'] + if config is None: + return self_attn.rope_theta + + rope_parameters = getattr(config, 'rope_parameters', None) + rope_type = _rope_type(config, rope_parameters if isinstance(rope_parameters, dict) else None) + if rope_type not in _UNSCALED_ROPE_TYPES: + raise ValueError(f"DeepSpeed kernel injection cannot serve rope_type={rope_type!r}. The injected " + "attention kernel builds its rotary embedding from rope_theta alone, so the " + "scaling parameters this configuration carries would be dropped and the model " + "would run with unscaled positions. Run this model without kernel injection " + "(replace_with_kernel_inject=False).") + + if hasattr(config, 'rope_theta'): + return config.rope_theta + if isinstance(rope_parameters, dict) and 'rope_theta' in rope_parameters: + return rope_parameters['rope_theta'] return self_attn.rope_theta diff --git a/tests/unit/module_inject/test_llama_rope_theta.py b/tests/unit/module_inject/test_llama_rope_theta.py index 73dd2ee88af1..0722d9e63678 100644 --- a/tests/unit/module_inject/test_llama_rope_theta.py +++ b/tests/unit/module_inject/test_llama_rope_theta.py @@ -89,3 +89,54 @@ def test_the_two_spellings_do_not_disagree_on_a_real_config(): def test_raises_when_nothing_carries_it(): with pytest.raises(AttributeError): _get_rope_theta(SimpleNamespace(config=SimpleNamespace())) + + +# --- scaled rotary variants ---------------------------------------------------- +# +# The injected kernel builds its rotary embedding from a scalar base +# (`InferenceContext.get_rotary(rotary_dim, rope_theta)`) and carries no scaling +# parameters at all, so a config asking for one cannot be served here. + + +@pytest.mark.parametrize("rope_type", ["llama3", "linear", "dynamic", "yarn", "longrope"]) +def test_a_scaled_rope_variant_is_refused(rope_type): + """Reading only rope_theta out of a scaled config is silently wrong. + + DeepSeek-R1-Distill-Llama-8B (#8340) is the live case: `rope_type="llama3"` with + `factor`, `low_freq_factor`, `high_freq_factor` and `original_max_position_embeddings`. + Dropping those and keeping the base runs the model with unscaled positions and no error, + which is worse than the AttributeError this helper exists to remove. + """ + config = SimpleNamespace(rope_parameters={ + "rope_type": rope_type, + "rope_theta": 500000.0, + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + }) + + with pytest.raises(ValueError, match="cannot serve rope_type"): + _get_rope_theta(SimpleNamespace(config=config)) + + +def test_a_scaled_variant_in_the_legacy_rope_scaling_spelling_is_refused(): + """transformers < 5.0 carries the same request under `rope_scaling`.""" + config = SimpleNamespace(rope_theta=500000.0, rope_scaling={"rope_type": "llama3", "factor": 8.0}) + + with pytest.raises(ValueError, match="cannot serve rope_type"): + _get_rope_theta(SimpleNamespace(config=config)) + + +def test_the_default_rope_type_is_not_refused(): + """`rope_type: "default"` is what standardize_rope_params writes for plain RoPE.""" + config = SimpleNamespace(rope_parameters={"rope_theta": 500000.0, "rope_type": "default"}) + + assert _get_rope_theta(SimpleNamespace(config=config)) == 500000.0 + + +def test_a_real_llama_config_is_not_refused(): + """The stock config the crash fix targets carries no scaling and must still resolve.""" + LlamaConfig = pytest.importorskip("transformers.models.llama.configuration_llama").LlamaConfig + + assert _get_rope_theta(SimpleNamespace(config=LlamaConfig(rope_theta=500000.0))) == 500000.0 From 4c687f12b98974b2ef086e4cb3cebe0e65fb29bc Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Tue, 8 Sep 2026 12:31:04 +0800 Subject: [PATCH 4/4] Apply yapf formatting to the rope_theta test The formatting job reformatted the SimpleNamespace call in test_a_scaled_rope_variant_is_refused. No behaviour change. Signed-off-by: alanhuangyoo --- .../unit/module_inject/test_llama_rope_theta.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/unit/module_inject/test_llama_rope_theta.py b/tests/unit/module_inject/test_llama_rope_theta.py index 0722d9e63678..6e040e37662c 100644 --- a/tests/unit/module_inject/test_llama_rope_theta.py +++ b/tests/unit/module_inject/test_llama_rope_theta.py @@ -107,14 +107,15 @@ def test_a_scaled_rope_variant_is_refused(rope_type): Dropping those and keeping the base runs the model with unscaled positions and no error, which is worse than the AttributeError this helper exists to remove. """ - config = SimpleNamespace(rope_parameters={ - "rope_type": rope_type, - "rope_theta": 500000.0, - "factor": 8.0, - "low_freq_factor": 1.0, - "high_freq_factor": 4.0, - "original_max_position_embeddings": 8192, - }) + config = SimpleNamespace( + rope_parameters={ + "rope_type": rope_type, + "rope_theta": 500000.0, + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + }) with pytest.raises(ValueError, match="cannot serve rope_type"): _get_rope_theta(SimpleNamespace(config=config))