feat: add test for positional rotary embeddings

server/tests/utils/test_rotary_emb.py

@@ -0,0 +1,281 @@
+import pytest
+import torch
+from unittest.mock import Mock, patch
+from text_generation_server.layers.rotary import (
+ PositionRotaryEmbedding,
+ DynamicPositionRotaryEmbedding,
+ YarnPositionRotaryEmbedding,
+)
+from text_generation_server.utils.import_utils import SYSTEM
+
+
+def test_position_rotary_embedding_static_basic():
+ config = Mock(rope_theta=10000, max_position_embeddings=2048, rope_scaling=None)
+ weights = Mock(device=torch.device("cpu"))
+
+ result = PositionRotaryEmbedding.static(
+ config=config,
+ dim=64,
+ base=config.rope_theta,
+ device=weights.device,
+ )
+
+ assert isinstance(result, PositionRotaryEmbedding)
+ assert result.inv_freq.shape == (32,) # dim // 2
+ assert result.scaling_factor is None
+
+
+def test_position_rotary_embedding_static_linear_scaling():
+ config = Mock(rope_theta=10000, max_position_embeddings=2048)
+ # scaling is not applied if type is linear (TODO: maybe revisit this)
+ config.rope_scaling = {"type": "linear", "factor": 2.0}
+ weights = Mock(device=torch.device("cpu"))
+
+ result = PositionRotaryEmbedding.static(
+ config=config,
+ dim=64,
+ base=config.rope_theta,
+ device=weights.device,
+ )
+
+ assert isinstance(result, PositionRotaryEmbedding)
+ assert result.scaling_factor is None
+
+
+def test_position_rotary_embedding_static_dynamic_scaling():
+ config = Mock(
+ rope_theta=10000,
+ max_position_embeddings=2048,
+ rope_scaling={"type": "dynamic", "factor": 2.0},
+ )
+ weights = Mock(device=torch.device("cpu"))
+
+ result = PositionRotaryEmbedding.static(
+ config=config,
+ dim=64,
+ base=config.rope_theta,
+ device=weights.device,
+ )
+
+ assert isinstance(result, DynamicPositionRotaryEmbedding)
+ assert result.scaling_factor == 2.0
+ assert result.max_position_embeddings == 2048
+
+
+def test_position_rotary_embedding_static_yarn_scaling():
+ config = Mock(
+ rope_theta=10000,
+ max_position_embeddings=2048,
+ rope_scaling={
+ "type": "yarn",
+ "factor": 1.5,
+ "original_max_position_embeddings": 2048,
+ },
+ )
+ weights = Mock(device=torch.device("cpu"))
+
+ result = PositionRotaryEmbedding.static(
+ config=config,
+ dim=64,
+ base=config.rope_theta,
+ device=weights.device,
+ )
+
+ assert isinstance(result, YarnPositionRotaryEmbedding)
+ assert result.scaling_factor == 1.5
+ assert result.max_position_embeddings == 2048
+
+
+def test_position_rotary_embedding_static_invalid_scaling():
+ config = Mock(
+ rope_theta=10000,
+ max_position_embeddings=2048,
+ rope_scaling={"type": "invalid", "factor": 2.0},
+ )
+ weights = Mock(device=torch.device("cpu"))
+
+ with pytest.raises(NotImplementedError):
+ PositionRotaryEmbedding.static(
+ config=config,
+ dim=64,
+ base=config.rope_theta,
+ device=weights.device,
+ )
+
+
+def test_position_rotary_embedding_static_llama3_scaling():
+ config = Mock(
+ rope_theta=10000,
+ max_position_embeddings=2048,
+ rope_scaling={
+ "rope_type": "llama3",
+ "factor": 2.0,
+ "low_freq_factor": 4,
+ "high_freq_factor": 32,
+ "original_max_position_embeddings": 2048,
+ },
+ )
+ weights = Mock(device=torch.device("cpu"))
+
+ result = PositionRotaryEmbedding.static(
+ config=config,
+ dim=64,
+ base=config.rope_theta,
+ device=weights.device,
+ )
+
+ assert isinstance(result, PositionRotaryEmbedding)
+ assert result.scaling_factor is None
+
+
+def test_position_rotary_embedding_max_tokens_exceed_max_position_embeddings():
+ config = Mock(
+ rope_theta=10000,
+ max_position_embeddings=4096,
+ rope_scaling=None,
+ )
+ weights = Mock(device=torch.device("cpu"))
+
+ with patch(
+ "text_generation_server.layers.rotary._get_rope_config"
+ ) as mock_get_rope_config:
+ mock_get_rope_config.return_value = {"type": "dynamic", "factor": 2.0}
+
+ result = PositionRotaryEmbedding.static(
+ config=config,
+ dim=64,
+ base=config.rope_theta,
+ device=weights.device,
+ )
+
+ assert isinstance(result, DynamicPositionRotaryEmbedding)
+ assert result.scaling_factor == 2.0
+ assert result.max_position_embeddings == 4096
+
+
+# Test the application of the rotary embedding
+
+
+def position_rotary_embedding_no_rope_config():
+ head_dim = 64
+ base = 10000
+ max_position_embeddings = 2048
+ num_heads = 16
+ batch_size = 2
+ seq_len = 128
+
+ device = "cuda"
+ dtype = torch.float16
+
+ config = Mock(
+ rope_theta=base,
+ max_position_embeddings=max_position_embeddings,
+ rope_scaling=None,
+ )
+
+ # create PositionRotaryEmbedding instance
+ rotary_emb = PositionRotaryEmbedding.static(
+ config=config, dim=head_dim, base=base, device=device
+ )
+
+ # generate position IDs
+ position_ids = torch.arange(seq_len).unsqueeze(0)
+ position_ids = position_ids.to(device).to(torch.int32).view(-1)
+
+ # get cos and sin values for the position IDs
+ cos, sin = rotary_emb.get_cos_sin(
+ position_ids=position_ids,
+ max_s=seq_len,
+ dtype=dtype,
+ )
+
+ # create query and key tensors
+ query = torch.randn(batch_size, seq_len, num_heads, head_dim).to(device).to(dtype)
+ key = torch.randn(batch_size, seq_len, num_heads, head_dim).to(device).to(dtype)
+
+ # clone to compare later
+ original_query = query.clone()
+ original_key = key.clone()
+
+ # apply rotary embedding
+ rotary_emb(query, key, cos, sin)
+
+ # copy rotated query and key and original query and key
+ q_rotated = query
+ k_rotated = key
+ query = original_query
+ key = original_key
+
+ assert (
+ q_rotated.shape == query.shape
+ ), "query shape should not change after rotation"
+ assert k_rotated.shape == key.shape, "key shape should not change after rotation"
+ assert not torch.allclose(q_rotated, query), "query should be modified by rotation"
+ assert not torch.allclose(k_rotated, key), "key should be modified by rotation"
+
+
+def position_rotary_embedding_with_dynamic_scaling():
+ head_dim = 64
+ base = 10000
+ max_position_embeddings = 2048
+ num_heads = 16
+ batch_size = 2
+ seq_len = 128
+
+ device = "cuda"
+ dtype = torch.float16
+
+ config = Mock(
+ rope_theta=base,
+ max_position_embeddings=max_position_embeddings,
+ rope_scaling={"type": "dynamic", "factor": 1.0},
+ )
+
+ # create PositionRotaryEmbedding instance
+ rotary_emb = PositionRotaryEmbedding.static(
+ config=config, dim=head_dim, base=base, device=device
+ )
+
+ # generate position IDs
+ position_ids = torch.arange(seq_len).unsqueeze(0)
+ position_ids = position_ids.to(device).to(torch.int32).view(-1)
+
+ # get cos and sin values for the position IDs
+ cos, sin = rotary_emb.get_cos_sin(
+ position_ids=position_ids,
+ max_s=seq_len,
+ dtype=dtype,
+ )
+
+ # create query and key tensors
+ query = torch.randn(batch_size, seq_len, num_heads, head_dim).to(device).to(dtype)
+ key = torch.randn(batch_size, seq_len, num_heads, head_dim).to(device).to(dtype)
+
+ # clone to compare later
+ original_query = query.clone()
+ original_key = key.clone()
+
+ # apply rotary embedding
+ rotary_emb(query, key, cos, sin)
+
+ # copy rotated query and key and original query and key
+ q_rotated = query
+ k_rotated = key
+ query = original_query
+ key = original_key
+
+ assert (
+ q_rotated.shape == query.shape
+ ), "query shape should not change after rotation"
+ assert k_rotated.shape == key.shape, "key shape should not change after rotation"
+ assert not torch.allclose(q_rotated, query), "query should be modified by rotation"
+ assert not torch.allclose(k_rotated, key), "key should be modified by rotation"
+
+
+if SYSTEM == "cuda":
+
+ def test_position_rotary_embedding_with_dynamic_scaling():
+ position_rotary_embedding_no_rope_config()
+
+ def test_position_rotary_embedding_no_rope_config():
+ position_rotary_embedding_no_rope_config()