diff --git a/deepspeed/inference/v2/model_implementations/exaone4/model.py b/deepspeed/inference/v2/model_implementations/exaone4/model.py index 5aeb87c2cdac..e12f5a513072 100644 --- a/deepspeed/inference/v2/model_implementations/exaone4/model.py +++ b/deepspeed/inference/v2/model_implementations/exaone4/model.py @@ -94,8 +94,7 @@ def positional_embedding_type(self) -> PositionalEmbeddingType: @property def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - rope_theta = getattr(self._config, "rope_theta", 1000000.0) - return RotateHalfConfig(theta_base=rope_theta) + return RotateHalfConfig(theta_base=self.rope_theta) def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/deepspeed/inference/v2/model_implementations/inference_transformer_base.py b/deepspeed/inference/v2/model_implementations/inference_transformer_base.py index fae67dc8fc2a..4971c072777e 100644 --- a/deepspeed/inference/v2/model_implementations/inference_transformer_base.py +++ b/deepspeed/inference/v2/model_implementations/inference_transformer_base.py @@ -45,6 +45,29 @@ def cached_property(func): return property(func) +# `RotateHalfConfig` carries `use_trained_freqs`, `theta_base` and `rotate_dim`, so a config +# asking for a scaled rotary variant cannot be honoured here. These spellings mean "no scaling". +_UNSCALED_ROPE_TYPES = (None, "default") + + +def _rope_types(config) -> set: + """Every rotary variant a config asks for, across both layouts and per-layer entries. + + transformers 5.x keeps these in ``rope_parameters``, 4.x in ``rope_scaling``, and a + per-layer-type config nests one dict per layer type inside either of them. Both + spellings of the key are in use, so both are read. + """ + rope_types = set() + for source in (getattr(config, "rope_parameters", None), getattr(config, "rope_scaling", None)): + if not isinstance(source, dict): + continue + for candidate in (source, *(value for value in source.values() if isinstance(value, dict))): + rope_type = candidate.get("rope_type", candidate.get("type")) + if rope_type is not None: + rope_types.add(rope_type) + return rope_types + + class DSTransformerModelBase(DSInferenceModelBase): """ Dimensioning properties @@ -165,6 +188,61 @@ def positional_embedding_config(self) -> Optional[RotateHalfConfig]: Derived helpers """ + @property + def rope_theta(self) -> float: + """The rotary base, read from wherever the installed transformers keeps it. + + transformers 5.0 folded the rotary settings into ``config.rope_parameters`` and + dropped the ``rope_theta`` attribute, so reading the attribute alone raises + against a stock config on 5.x. + + A config that sets RoPE per layer type gets nested one level deeper, keyed by + the layer type, and ``standardize_rope_params`` leaves the class default at the + top level of the same dict. Reading the top level there returns that default + rather than anything the checkpoint asked for, so the nested entries win. Every + caller of this property feeds a single ``RotateHalfConfig.theta_base`` for the + whole model, so distinct per-layer bases cannot be represented and are refused + rather than silently resolved to one of them. + + A scaled variant is refused for the same reason, matching what #8341 does for + kernel injection: every caller builds ``RotateHalfConfig(theta_base=...)``, and + that config carries ``use_trained_freqs``, ``theta_base`` and ``rotate_dim`` and + nothing else, so the scaling parameters have nowhere to go. Returning the base + alone would run the model with unscaled positions and no error. + """ + scaled = sorted(rope_type for rope_type in _rope_types(self._config) + if rope_type not in _UNSCALED_ROPE_TYPES) + if scaled: + raise ValueError(f"Inference V2 cannot serve rope_type={scaled[0]!r} " + f"({type(self._config).__name__}). The rotary embedding is built from " + "theta_base alone, so the scaling parameters this configuration carries " + "would be dropped and the model would run with unscaled positions.") + + theta = getattr(self._config, "rope_theta", None) + if theta is not None: + return theta + + rope_parameters = getattr(self._config, "rope_parameters", None) or getattr(self._config, "rope_scaling", None) + if isinstance(rope_parameters, dict): + per_layer = { + layer_type: parameters["rope_theta"] + for layer_type, parameters in rope_parameters.items() + if isinstance(parameters, dict) and parameters.get("rope_theta") is not None + } + if per_layer: + distinct = set(per_layer.values()) + if len(distinct) > 1: + raise ValueError(f"{type(self._config).__name__} sets a different rope_theta per " + f"layer type ({per_layer}); Inference V2 applies one rotary base " + "to every layer and cannot represent this config.") + return distinct.pop() + + if rope_parameters.get("rope_theta") is not None: + return rope_parameters["rope_theta"] + + raise ValueError(f"{type(self._config).__name__} carries no rope_theta, either as an " + "attribute or in rope_parameters/rope_scaling.") + @cached_property def n_heads_q_local(self) -> int: """ diff --git a/deepspeed/inference/v2/model_implementations/llama_v2/model.py b/deepspeed/inference/v2/model_implementations/llama_v2/model.py index a0c81f4d749e..ac23145e942e 100644 --- a/deepspeed/inference/v2/model_implementations/llama_v2/model.py +++ b/deepspeed/inference/v2/model_implementations/llama_v2/model.py @@ -107,7 +107,7 @@ def positional_embedding_type(self) -> PositionalEmbeddingType: @property def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) + return RotateHalfConfig(theta_base=self.rope_theta) """ Forward implementations diff --git a/deepspeed/inference/v2/model_implementations/mistral/model.py b/deepspeed/inference/v2/model_implementations/mistral/model.py index 318d362f1a64..89289a9f171a 100644 --- a/deepspeed/inference/v2/model_implementations/mistral/model.py +++ b/deepspeed/inference/v2/model_implementations/mistral/model.py @@ -106,7 +106,7 @@ def positional_embedding_type(self) -> PositionalEmbeddingType: @property def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) + return RotateHalfConfig(theta_base=self.rope_theta) """ Forward implementations diff --git a/deepspeed/inference/v2/model_implementations/mixtral/model.py b/deepspeed/inference/v2/model_implementations/mixtral/model.py index 878cd8e31cec..5a512fc0633a 100644 --- a/deepspeed/inference/v2/model_implementations/mixtral/model.py +++ b/deepspeed/inference/v2/model_implementations/mixtral/model.py @@ -114,7 +114,7 @@ def positional_embedding_config(self) -> Optional[RotateHalfConfig]: """ The positional embedding configuration for the model. """ - return RotateHalfConfig(theta_base=self._config.rope_theta) + return RotateHalfConfig(theta_base=self.rope_theta) """ Inherited from `DSMoETransformerModelBase` diff --git a/deepspeed/inference/v2/model_implementations/phi/model.py b/deepspeed/inference/v2/model_implementations/phi/model.py index 2d5826810cb5..41e0e8652d60 100644 --- a/deepspeed/inference/v2/model_implementations/phi/model.py +++ b/deepspeed/inference/v2/model_implementations/phi/model.py @@ -97,7 +97,7 @@ def positional_embedding_type(self) -> PositionalEmbeddingType: @property def positional_embedding_config(self) -> Optional[RotateHalfConfig]: rotary_dim = int(self._config.partial_rotary_factor * self.head_size) - return RotateHalfConfig(rotate_dim=rotary_dim, theta_base=self._config.rope_theta) + return RotateHalfConfig(rotate_dim=rotary_dim, theta_base=self.rope_theta) """ Forward implementations diff --git a/deepspeed/inference/v2/model_implementations/phi3/model.py b/deepspeed/inference/v2/model_implementations/phi3/model.py index 507bb4fc9af1..7d974d3a113e 100644 --- a/deepspeed/inference/v2/model_implementations/phi3/model.py +++ b/deepspeed/inference/v2/model_implementations/phi3/model.py @@ -106,7 +106,7 @@ def positional_embedding_type(self) -> PositionalEmbeddingType: @property def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) + return RotateHalfConfig(theta_base=self.rope_theta) """ Forward implementations diff --git a/deepspeed/inference/v2/model_implementations/qwen_v2/model.py b/deepspeed/inference/v2/model_implementations/qwen_v2/model.py index d535462a954d..c89ed8a65c65 100644 --- a/deepspeed/inference/v2/model_implementations/qwen_v2/model.py +++ b/deepspeed/inference/v2/model_implementations/qwen_v2/model.py @@ -100,7 +100,7 @@ def positional_embedding_type(self) -> PositionalEmbeddingType: @property def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) + return RotateHalfConfig(theta_base=self.rope_theta) def make_norm_layer(self) -> None: """ diff --git a/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py b/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py index c7841b24e5fc..bc441452ce70 100644 --- a/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py +++ b/deepspeed/inference/v2/model_implementations/qwen_v2_moe/model.py @@ -105,7 +105,7 @@ def positional_embedding_type(self) -> PositionalEmbeddingType: @property def positional_embedding_config(self) -> Optional[RotateHalfConfig]: - return RotateHalfConfig(theta_base=self._config.rope_theta) + return RotateHalfConfig(theta_base=self.rope_theta) """ Inherited from `DSMoETransformerModelBase` diff --git a/tests/unit/inference/v2/model_implementations/test_rope_theta.py b/tests/unit/inference/v2/model_implementations/test_rope_theta.py new file mode 100644 index 000000000000..e74170e525f2 --- /dev/null +++ b/tests/unit/inference/v2/model_implementations/test_rope_theta.py @@ -0,0 +1,194 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Inference V2 must find rope_theta wherever the installed transformers keeps it. + +transformers 5.0 folded the rotary settings into ``config.rope_parameters`` and dropped +the ``rope_theta`` attribute. Models that read the attribute directly raise against a +stock config on 5.x; ``exaone4`` read it through a ``getattr`` default and silently used +that default instead of the value the config actually carries. +""" + +from types import SimpleNamespace + +import pytest + +from deepspeed.inference.v2.model_implementations.inference_transformer_base import DSTransformerModelBase + + +class _Model: + """Minimal stand-in that borrows the property under test.""" + + rope_theta = DSTransformerModelBase.rope_theta + + def __init__(self, config): + self._config = config + + +def test_reads_the_legacy_attribute(): + assert _Model(SimpleNamespace(rope_theta=500000.0)).rope_theta == 500000.0 + + +def test_reads_rope_parameters_when_the_attribute_is_gone(): + config = SimpleNamespace(rope_parameters={"rope_theta": 10000.0, "rope_type": "default"}) + + assert _Model(config).rope_theta == 10000.0 + + +def test_reads_rope_scaling_when_that_is_the_only_dict(): + config = SimpleNamespace(rope_scaling={"rope_theta": 1000000.0, "rope_type": "default"}) + + assert _Model(config).rope_theta == 1000000.0 + + +def test_prefers_the_attribute_when_both_are_present(): + config = SimpleNamespace(rope_theta=500000.0, rope_parameters={"rope_theta": 10000.0}) + + assert _Model(config).rope_theta == 500000.0 + + +def test_raises_instead_of_guessing_when_nothing_carries_it(): + # exaone4 used to fall back to 1e6 here, which is a silent 100x error against a + # config whose real base is 1e4. + with pytest.raises(ValueError, match="rope_theta"): + _ = _Model(SimpleNamespace()).rope_theta + + +def test_unwraps_a_per_layer_type_dict(): + """A config that sets RoPE per layer type nests one level deeper. + + ``standardize_rope_params`` leaves the class default at the top level of the same + dict, so reading the top level returns 10000.0 rather than the 1000000.0 the + checkpoint asked for. + """ + config = SimpleNamespace( + rope_parameters={ + "sliding_attention": { + "rope_theta": 1000000.0, + "rope_type": "default" + }, + "full_attention": { + "rope_theta": 1000000.0, + "rope_type": "default" + }, + "rope_theta": 10000.0, + "rope_type": "default", + }) + + assert _Model(config).rope_theta == 1000000.0 + + +def test_refuses_a_config_whose_layer_types_disagree(): + """Callers feed one ``RotateHalfConfig.theta_base`` for the whole model. + + Picking either base would be wrong for the layers using the other one, so this is + refused rather than resolved. + """ + config = SimpleNamespace( + rope_parameters={ + "sliding_attention": { + "rope_theta": 1000000.0, + "rope_type": "default" + }, + "full_attention": { + "rope_theta": 16000000.0, + "rope_type": "default" + }, + "rope_theta": 10000.0, + "rope_type": "default", + }) + + with pytest.raises(ValueError, match="different rope_theta per layer type"): + _ = _Model(config).rope_theta + + +@pytest.mark.parametrize( + "module_name, config_name", + [ + ("llama", "LlamaConfig"), + ("mistral", "MistralConfig"), + ("mixtral", "MixtralConfig"), + ("phi", "PhiConfig"), + ("phi3", "Phi3Config"), + ("qwen2", "Qwen2Config"), + ("qwen2_moe", "Qwen2MoeConfig"), + ("exaone4", "Exaone4Config"), + ], +) +def test_resolves_against_the_installed_transformers_configs(module_name, config_name): + """Every config backing a V2 model must yield a base through one spelling or the other.""" + importlib = pytest.importorskip("importlib") + try: + module = importlib.import_module(f"transformers.models.{module_name}.configuration_{module_name}") + config = getattr(module, config_name)() + except (ImportError, AttributeError): + pytest.skip(f"{config_name} is not available in the installed transformers") + + assert _Model(config).rope_theta > 0 + + +# --- scaled variants ------------------------------------------------------------- +# +# Every caller builds `RotateHalfConfig(theta_base=self.rope_theta)`, and that config +# carries `use_trained_freqs`, `theta_base` and `rotate_dim` and nothing else. A config +# asking for a scaled rotary has nowhere to put its scaling parameters, so returning the +# base alone would run the model with unscaled positions and no error. #8341 made kernel +# injection refuse the same thing; this is the Inference V2 half. + + +@pytest.mark.parametrize("rope_type", ["llama3", "linear", "dynamic", "yarn", "longrope"]) +def test_rejects_a_scaled_variant_in_rope_parameters(rope_type): + config = SimpleNamespace(rope_parameters={"rope_theta": 500000.0, "rope_type": rope_type}) + + with pytest.raises(ValueError, match=rope_type): + _Model(config).rope_theta + + +@pytest.mark.parametrize("key", ["rope_type", "type"]) +def test_rejects_a_scaled_variant_in_legacy_rope_scaling(key): + """transformers 4.x keeps this in `rope_scaling`, and both spellings of the key are in use.""" + config = SimpleNamespace(rope_theta=500000.0, rope_scaling={key: "llama3", "factor": 8.0}) + + with pytest.raises(ValueError, match="llama3"): + _Model(config).rope_theta + + +def test_rejects_a_scaled_variant_nested_per_layer_type(): + config = SimpleNamespace( + rope_parameters={ + "sliding_attention": { + "rope_theta": 1000000.0, + "rope_type": "default" + }, + "full_attention": { + "rope_theta": 1000000.0, + "rope_type": "yarn", + "factor": 4.0 + }, + "rope_theta": 10000.0, + "rope_type": "default", + }) + + with pytest.raises(ValueError, match="yarn"): + _Model(config).rope_theta + + +def test_a_scaled_variant_is_refused_even_when_the_attribute_carries_the_base(): + """The attribute read comes first, so the check has to precede it to be reachable.""" + config = SimpleNamespace(rope_theta=500000.0, rope_parameters={"rope_type": "llama3", "factor": 8.0}) + + with pytest.raises(ValueError, match="llama3"): + _Model(config).rope_theta + + +@pytest.mark.parametrize("rope_type", [None, "default"]) +def test_unscaled_configurations_still_resolve(rope_type): + """The two spellings that mean "no scaling" must keep working, in both layouts.""" + parameters = {"rope_theta": 500000.0} + if rope_type is not None: + parameters["rope_type"] = rope_type + + assert _Model(SimpleNamespace(rope_parameters=dict(parameters))).rope_theta == 500000.0 + assert _Model(SimpleNamespace(rope_scaling=dict(parameters))).rope_theta == 500000.0 +