Add files using upload-large-folder tool

babe29f verified 24 days ago

6.41 kB

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import torch

from models.tt_transformers.tt.common import gather_cos_sin, precompute_freqs, rope_scaling_model_factory
from models.tt_transformers.tt.rope import RotaryEmbedding, rotary_embedding_factory


class TestRope:
    """Test suite to compare different RoPE implementations for consistency."""

    def test_basic_rope_vs_precompute_freqs(self):
        """
        Test that compares sin/cos matrices computed by RotaryEmbedding class
        vs precompute_freqs function to check for discrepancies.
        """
        # Test parameters
        dim = 128
        max_seq_len = 1024
        base = 10000.0
        device = torch.device("cpu")

        # Create RotaryEmbedding instance
        rope = RotaryEmbedding(dim=dim, max_position_embeddings=max_seq_len, base=base, device=device)

        # Get cos/sin from RotaryEmbedding
        rope_cos, rope_sin = rope.cos_cached, rope.sin_cached

        # Get cos/sin from precompute_freqs
        precompute_cos, precompute_sin = precompute_freqs(
            dim=dim, end=2 * max_seq_len, theta=base, scale_factor=None, orig_context_len=None
        )
        precompute_cos, precompute_sin = gather_cos_sin(torch.arange(max_seq_len), precompute_cos, precompute_sin)

        print(f"RotaryEmbedding cos shape: {rope_cos.shape}")
        print(f"RotaryEmbedding sin shape: {rope_sin.shape}")
        print(f"precompute_freqs cos shape: {precompute_cos.shape}")
        print(f"precompute_freqs sin shape: {precompute_sin.shape}")

        # Compare shapes
        assert (
            rope_cos.shape == precompute_cos.shape
        ), f"Cos shapes don't match: {rope_cos.shape} vs {precompute_cos.shape}"
        assert (
            rope_sin.shape == precompute_sin.shape
        ), f"Sin shapes don't match: {rope_sin.shape} vs {precompute_sin.shape}"

        # Compare values with tolerance
        cos_diff = torch.abs(rope_cos - precompute_cos)
        sin_diff = torch.abs(rope_sin - precompute_sin)

        max_cos_diff = torch.max(cos_diff)
        max_sin_diff = torch.max(sin_diff)

        print(f"Max cos difference: {max_cos_diff}")
        print(f"Max sin difference: {max_sin_diff}")
        print(f"Mean cos difference: {torch.mean(cos_diff)}")
        print(f"Mean sin difference: {torch.mean(sin_diff)}")

        # Allow for small numerical differences
        tolerance = 1e-6
        assert max_cos_diff < tolerance, f"Cos values differ by more than {tolerance}: {max_cos_diff}"
        assert max_sin_diff < tolerance, f"Sin values differ by more than {tolerance}: {max_sin_diff}"

    def test_rope_llama3_scaling(self):
        """
        Test that the shape of the cos/sin matrices is correct for yarn scaling.
        """
        dim = 128
        max_seq_len = 1024
        base = 10000.0
        device = torch.device("cpu")

        rope = RotaryEmbedding(dim=dim, max_position_embeddings=max_seq_len, base=base, device=device)
        rope_cos, rope_sin = rope.cos_cached, rope.sin_cached

        rope_llama_model = rope_scaling_model_factory(
            {"rope_type": "llama3", "factor": 32, "original_max_position_embeddings": 8192}
        )
        rope_llama_scaled = rotary_embedding_factory(
            dim=dim, max_position_embeddings=max_seq_len, base=base, rope_scaling=rope_llama_model
        )
        rope_llama_scaled_cos, rope_llama_scaled_sin = rope_llama_scaled.cos_cached, rope_llama_scaled.sin_cached

        assert rope_llama_scaled_cos.shape == rope_cos.shape == (1, 1, max_seq_len, dim)
        assert rope_llama_scaled_sin.shape == rope_sin.shape == (1, 1, max_seq_len, dim)

        cos_diff = torch.abs(rope_cos - rope_llama_scaled_cos)
        sin_diff = torch.abs(rope_sin - rope_llama_scaled_sin)

        max_cos_diff = torch.max(cos_diff)
        max_sin_diff = torch.max(sin_diff)

        print(f"Max cos difference: {max_cos_diff}")
        print(f"Max sin difference: {max_sin_diff}")
        print(f"Mean cos difference: {torch.mean(cos_diff)}")
        print(f"Mean sin difference: {torch.mean(sin_diff)}")

        # Make sure we actually ran the scaling
        assert max_cos_diff > 1e-6, f"Cos values are the same as non scaled. Max diff = {max_cos_diff}"
        assert max_sin_diff > 1e-6, f"Sin values are the same as non scaled. Max diff = {max_sin_diff}"

    def test_rope_yarn_scaling(self):
        """
        Test that the shape of the cos/sin matrices is correct for yarn scaling.
        """
        dim = 128
        max_seq_len = 1024
        base = 10000.0
        device = torch.device("cpu")

        rope = RotaryEmbedding(dim=dim, max_position_embeddings=max_seq_len, base=base, device=device)
        rope_cos, rope_sin = rope.cos_cached, rope.sin_cached

        rope_yarn_model = rope_scaling_model_factory(
            {"rope_type": "yarn", "factor": 32, "original_max_position_embeddings": 8192}
        )
        rope_yarn_scaled = rotary_embedding_factory(
            dim=dim, max_position_embeddings=max_seq_len, base=base, rope_scaling=rope_yarn_model
        )
        rope_yarn_scaled_cos, rope_yarn_scaled_sin = rope_yarn_scaled.cos_cached, rope_yarn_scaled.sin_cached

        assert rope_yarn_scaled_cos.shape == rope_cos.shape == (1, 1, max_seq_len, dim)
        assert rope_yarn_scaled_sin.shape == rope_sin.shape == (1, 1, max_seq_len, dim)

        cos_diff = torch.abs(rope_cos - rope_yarn_scaled_cos)
        sin_diff = torch.abs(rope_sin - rope_yarn_scaled_sin)

        max_cos_diff = torch.max(cos_diff)
        max_sin_diff = torch.max(sin_diff)

        print(f"Max cos difference: {max_cos_diff}")
        print(f"Max sin difference: {max_sin_diff}")
        print(f"Mean cos difference: {torch.mean(cos_diff)}")
        print(f"Mean sin difference: {torch.mean(sin_diff)}")

        # Make sure we actually ran the scaling
        assert max_cos_diff > 1e-6, f"Cos values are the same as non scaled. Max diff = {max_cos_diff}"
        assert max_sin_diff > 1e-6, f"Sin values are the same as non scaled. Max diff = {max_sin_diff}"


if __name__ == "__main__":
    # Run a quick test if executed directly
    test_instance = TestRope()
    test_instance.test_basic_rope_vs_precompute_freqs()
    test_instance.test_rope_llama3_scaling_shape()
    test_instance.test_rope_yarn_scaling_shape()
    print("All tests passed!")

	# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

	# SPDX-License-Identifier: Apache-2.0

	import torch

	from models.tt_transformers.tt.common import gather_cos_sin, precompute_freqs, rope_scaling_model_factory
	from models.tt_transformers.tt.rope import RotaryEmbedding, rotary_embedding_factory


	class TestRope:
	"""Test suite to compare different RoPE implementations for consistency."""

	def test_basic_rope_vs_precompute_freqs(self):
	"""
	Test that compares sin/cos matrices computed by RotaryEmbedding class
	vs precompute_freqs function to check for discrepancies.
	"""
	# Test parameters
	dim = 128
	max_seq_len = 1024
	base = 10000.0
	device = torch.device("cpu")

	# Create RotaryEmbedding instance
	rope = RotaryEmbedding(dim=dim, max_position_embeddings=max_seq_len, base=base, device=device)

	# Get cos/sin from RotaryEmbedding
	rope_cos, rope_sin = rope.cos_cached, rope.sin_cached

	# Get cos/sin from precompute_freqs
	precompute_cos, precompute_sin = precompute_freqs(
	dim=dim, end=2 * max_seq_len, theta=base, scale_factor=None, orig_context_len=None
	)
	precompute_cos, precompute_sin = gather_cos_sin(torch.arange(max_seq_len), precompute_cos, precompute_sin)

	print(f"RotaryEmbedding cos shape: {rope_cos.shape}")
	print(f"RotaryEmbedding sin shape: {rope_sin.shape}")
	print(f"precompute_freqs cos shape: {precompute_cos.shape}")
	print(f"precompute_freqs sin shape: {precompute_sin.shape}")

	# Compare shapes
	assert (
	rope_cos.shape == precompute_cos.shape
	), f"Cos shapes don't match: {rope_cos.shape} vs {precompute_cos.shape}"
	assert (
	rope_sin.shape == precompute_sin.shape
	), f"Sin shapes don't match: {rope_sin.shape} vs {precompute_sin.shape}"

	# Compare values with tolerance
	cos_diff = torch.abs(rope_cos - precompute_cos)
	sin_diff = torch.abs(rope_sin - precompute_sin)

	max_cos_diff = torch.max(cos_diff)
	max_sin_diff = torch.max(sin_diff)

	print(f"Max cos difference: {max_cos_diff}")
	print(f"Max sin difference: {max_sin_diff}")
	print(f"Mean cos difference: {torch.mean(cos_diff)}")
	print(f"Mean sin difference: {torch.mean(sin_diff)}")

	# Allow for small numerical differences
	tolerance = 1e-6
	assert max_cos_diff < tolerance, f"Cos values differ by more than {tolerance}: {max_cos_diff}"
	assert max_sin_diff < tolerance, f"Sin values differ by more than {tolerance}: {max_sin_diff}"

	def test_rope_llama3_scaling(self):
	"""
	Test that the shape of the cos/sin matrices is correct for yarn scaling.
	"""
	dim = 128
	max_seq_len = 1024
	base = 10000.0
	device = torch.device("cpu")

	rope = RotaryEmbedding(dim=dim, max_position_embeddings=max_seq_len, base=base, device=device)
	rope_cos, rope_sin = rope.cos_cached, rope.sin_cached

	rope_llama_model = rope_scaling_model_factory(
	{"rope_type": "llama3", "factor": 32, "original_max_position_embeddings": 8192}
	)
	rope_llama_scaled = rotary_embedding_factory(
	dim=dim, max_position_embeddings=max_seq_len, base=base, rope_scaling=rope_llama_model
	)
	rope_llama_scaled_cos, rope_llama_scaled_sin = rope_llama_scaled.cos_cached, rope_llama_scaled.sin_cached

	assert rope_llama_scaled_cos.shape == rope_cos.shape == (1, 1, max_seq_len, dim)
	assert rope_llama_scaled_sin.shape == rope_sin.shape == (1, 1, max_seq_len, dim)

	cos_diff = torch.abs(rope_cos - rope_llama_scaled_cos)
	sin_diff = torch.abs(rope_sin - rope_llama_scaled_sin)

	max_cos_diff = torch.max(cos_diff)
	max_sin_diff = torch.max(sin_diff)

	print(f"Max cos difference: {max_cos_diff}")
	print(f"Max sin difference: {max_sin_diff}")
	print(f"Mean cos difference: {torch.mean(cos_diff)}")
	print(f"Mean sin difference: {torch.mean(sin_diff)}")

	# Make sure we actually ran the scaling
	assert max_cos_diff > 1e-6, f"Cos values are the same as non scaled. Max diff = {max_cos_diff}"
	assert max_sin_diff > 1e-6, f"Sin values are the same as non scaled. Max diff = {max_sin_diff}"

	def test_rope_yarn_scaling(self):
	"""
	Test that the shape of the cos/sin matrices is correct for yarn scaling.
	"""
	dim = 128
	max_seq_len = 1024
	base = 10000.0
	device = torch.device("cpu")

	rope = RotaryEmbedding(dim=dim, max_position_embeddings=max_seq_len, base=base, device=device)
	rope_cos, rope_sin = rope.cos_cached, rope.sin_cached

	rope_yarn_model = rope_scaling_model_factory(
	{"rope_type": "yarn", "factor": 32, "original_max_position_embeddings": 8192}
	)
	rope_yarn_scaled = rotary_embedding_factory(
	dim=dim, max_position_embeddings=max_seq_len, base=base, rope_scaling=rope_yarn_model
	)
	rope_yarn_scaled_cos, rope_yarn_scaled_sin = rope_yarn_scaled.cos_cached, rope_yarn_scaled.sin_cached

	assert rope_yarn_scaled_cos.shape == rope_cos.shape == (1, 1, max_seq_len, dim)
	assert rope_yarn_scaled_sin.shape == rope_sin.shape == (1, 1, max_seq_len, dim)

	cos_diff = torch.abs(rope_cos - rope_yarn_scaled_cos)
	sin_diff = torch.abs(rope_sin - rope_yarn_scaled_sin)

	max_cos_diff = torch.max(cos_diff)
	max_sin_diff = torch.max(sin_diff)

	print(f"Max cos difference: {max_cos_diff}")
	print(f"Max sin difference: {max_sin_diff}")
	print(f"Mean cos difference: {torch.mean(cos_diff)}")
	print(f"Mean sin difference: {torch.mean(sin_diff)}")

	# Make sure we actually ran the scaling
	assert max_cos_diff > 1e-6, f"Cos values are the same as non scaled. Max diff = {max_cos_diff}"
	assert max_sin_diff > 1e-6, f"Sin values are the same as non scaled. Max diff = {max_sin_diff}"


	if __name__ == "__main__":
	# Run a quick test if executed directly
	test_instance = TestRope()
	test_instance.test_basic_rope_vs_precompute_freqs()
	test_instance.test_rope_llama3_scaling_shape()
	test_instance.test_rope_yarn_scaling_shape()
	print("All tests passed!")