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 <alanhuangyoo@gmail.com>

Author: alanhuangyoo

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
 
 

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