diff --git a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py index d745df678e93..a199f49d8528 100644 --- a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py +++ b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py @@ -80,8 +80,16 @@ def softmax_context_fallback(self, query_key_value, attn_mask, rotary_dim, rotat from transformers.models.llama.modeling_llama import apply_rotary_pos_emb rotary = InferenceContext.Instance().get_rotary(rotary_dim, rope_theta, bat_0213_value.device) - cos, sin = rotary(bat_0213_value, InferenceContext.Instance().get_max_tokens_num()) - bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin, position_ids) + # LlamaRotaryEmbedding.forward takes position_ids, not a token count. + cos, sin = rotary(bat_0213_value, position_ids) + # apply_rotary_pos_emb takes them only through cos/sin. transformers 5.0 + # dropped the deprecated position_ids parameter, so the fifth positional + # slot is unsqueeze_dim there: + # 4.51.3 .. 4.57.0 (q, k, cos, sin, position_ids=None, unsqueeze_dim=1) + # 5.0.0 .. 5.16.1 (q, k, cos, sin, unsqueeze_dim=1) + # Passing four arguments is correct on both, and leaves unsqueeze_dim at its + # default rather than handing it a tensor. + bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin) bat_0213_key, bat_0213_value = InferenceContext.Instance().update_cache(layer_id, token_idx, is_prompt, bat_0213_key, bat_0213_value) diff --git a/deepspeed/ops/transformer/inference/op_binding/workspace.py b/deepspeed/ops/transformer/inference/op_binding/workspace.py index 0565666559bf..efd9104a57b4 100644 --- a/deepspeed/ops/transformer/inference/op_binding/workspace.py +++ b/deepspeed/ops/transformer/inference/op_binding/workspace.py @@ -137,9 +137,17 @@ def get_kv_cache(self): def get_rotary(self, rotary_dim, rope_theta, device=None): if self.rotary is None: + from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding - self.rotary = LlamaRotaryEmbedding(rotary_dim, base=rope_theta, device=device) + # transformers 4.48 replaced the (dim, base=..., device=...) signature with one + # that reads both off a config. head_dim carries the rotary width and rope_theta + # the base; transformers 5 folds rope_theta into rope_parameters itself, so + # passing it this way works on both. + config = LlamaConfig(head_dim=rotary_dim, rope_theta=rope_theta) + self.rotary = LlamaRotaryEmbedding(config) + if device is not None: + self.rotary = self.rotary.to(device) return self.rotary diff --git a/tests/unit/ops/transformer/inference/test_rotary_embedding.py b/tests/unit/ops/transformer/inference/test_rotary_embedding.py new file mode 100644 index 000000000000..f728e60fabfd --- /dev/null +++ b/tests/unit/ops/transformer/inference/test_rotary_embedding.py @@ -0,0 +1,112 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""InferenceContext.get_rotary has to build a rotary embedding the installed transformers accepts. + +transformers 4.48 replaced `LlamaRotaryEmbedding(dim, base=..., device=...)` with a +config-taking constructor, and its forward went from a token count to `position_ids`. The +fallback attention path in `softmax_context.py` used both of the old shapes, so it raised +`TypeError: __init__() got an unexpected keyword argument 'base'` before reaching any kernel. + +No accelerator needed: this covers the construction and the rope values, and the rest of the +fallback path is unchanged. +""" + +import pytest +import torch + +from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig +from deepspeed.ops.transformer.inference.op_binding.softmax_context import SoftmaxContextOp +from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext + + +def _reference_cos_sin(rotary_dim, rope_theta, seq_len): + """cos/sin straight from the rope definition, independent of transformers.""" + inv_freq = 1.0 / (rope_theta**(torch.arange(0, rotary_dim, 2).float() / rotary_dim)) + freqs = torch.outer(torch.arange(seq_len).float(), inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos(), emb.sin() + + +@pytest.fixture +def context(): + ctx = InferenceContext.Instance() + ctx.rotary = None + yield ctx + ctx.rotary = None + + +@pytest.mark.parametrize("rotary_dim", [32, 64, 128]) +@pytest.mark.parametrize("rope_theta", [10000.0, 500000.0]) +def test_get_rotary_matches_the_rope_definition(context, rotary_dim, rope_theta): + seq_len = 12 + rotary = context.get_rotary(rotary_dim, rope_theta) + + position_ids = torch.arange(seq_len).unsqueeze(0) + cos, sin = rotary(torch.zeros(1, 1, seq_len, rotary_dim), position_ids) + + expected_cos, expected_sin = _reference_cos_sin(rotary_dim, rope_theta, seq_len) + assert cos.shape == (1, seq_len, rotary_dim) + torch.testing.assert_close(cos[0], expected_cos) + torch.testing.assert_close(sin[0], expected_sin) + + +def test_get_rotary_uses_rotary_dim_not_the_config_default(context): + """rotary_dim has to reach the embedding, not the LlamaConfig hidden_size // heads default.""" + rotary = context.get_rotary(32, 10000.0) + + assert rotary.inv_freq.numel() == 16 + + +def test_get_rotary_is_cached(context): + first = context.get_rotary(64, 10000.0) + + assert context.get_rotary(64, 10000.0) is first + + +class _StopAfterRotary(Exception): + """Ends the fallback at the call under test, before it wants a workspace.""" + + +def test_the_fallback_passes_apply_rotary_pos_emb_four_arguments(monkeypatch, context): + """Asserts this repo's call, not the library's signature. + + transformers 5.0 dropped the deprecated `position_ids` parameter, so the fifth + positional slot became `unsqueeze_dim`: + + 4.51.3 .. 4.57.0 (q, k, cos, sin, position_ids=None, unsqueeze_dim=1) + 5.0.0 .. 5.16.1 (q, k, cos, sin, unsqueeze_dim=1) + + Passing `position_ids` there reaches `unsqueeze(dim=...)` as a tensor. Checking the + library's own signature would not catch that, since the mistake is in the caller; the + recorded call has to come from `softmax_context_fallback` itself. + + The recorder raises so execution stops at the rotary block. `update_cache` sits two + lines below and needs a workspace, which is not what this is about. + """ + llama = pytest.importorskip("transformers.models.llama.modeling_llama") + + recorded = {} + + def recorder(*args, **kwargs): + recorded["args"], recorded["kwargs"] = args, kwargs + raise _StopAfterRotary + + monkeypatch.setattr(llama, "apply_rotary_pos_emb", recorder) + + heads, head_dim, seq_len, rotary_dim = 4, 16, 8, 16 + query_key_value = torch.randn(1, seq_len, 3 * heads * head_dim) + position_ids = torch.arange(seq_len).unsqueeze(0) + + config = DeepSpeedInferenceConfig(hidden_size=heads * head_dim, heads=heads, dtype=torch.float32) + op = SoftmaxContextOp.__new__(SoftmaxContextOp) + op.config = config + + with pytest.raises(_StopAfterRotary): + op.softmax_context_fallback(query_key_value, None, rotary_dim, True, False, heads, heads, 1.0, False, False, 0, + False, 0, 1, None, 10000.0, True, 0, position_ids) + + assert len(recorded["args"]) == 4, \ + f"the fallback passed {len(recorded['args'])} positional arguments; the fifth is unsqueeze_dim" + assert not recorded["kwargs"], f"unexpected keyword arguments: {sorted(recorded['kwargs'])}"