Add unit tests to adopt tt_ccl

Author: nikileshx

models/tt_transformers/tests/multimodal/mistral_24b/test_patch_rot_emb.py

@@ -8,9 +8,10 @@
 from loguru import logger
 
 import ttnn
+from models.tt_transformers.tt.model_config import ModelArgs
 
+# models/tt_transformers/tt/common.py
 from models.tt_transformers.tt.multimodal.mistral_24b.vision_rope import VisionRotarySetup as RotarySetup
-from models.tt_transformers.tt.model_config import ModelArgs
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 
 

models/tt_transformers/tests/multimodal/mistral_24b/test_pixtral_transformer.py

@@ -8,8 +8,9 @@
 from loguru import logger
 
 import ttnn
-from models.experimental.mistral_24b.tt.vision_pixtral_transformer import TtPixtralTransformer
+from models.tt_transformers.tt.ccl import TT_CCL
 from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.multimodal.mistral_24b.vision_pixtral_transformer import TtPixtralTransformer
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 
 
@@ -27,6 +28,11 @@
     ],
     indirect=True,
 )
+@pytest.mark.parametrize(
+    "device_params",
+    [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "trace_region_size": 30000000, "num_command_queues": 1}],
+    indirect=True,
+)
 def test_image_transformer_inference(batch, num_chunks, mesh_device):
     pcc_required = 0.99
 
@@ -51,8 +57,10 @@ def test_image_transformer_inference(batch, num_chunks, mesh_device):
 
     all_tests_pass = True
 
+    tt_ccl = TT_CCL(mesh_device)
     tt_model = TtPixtralTransformer(
         mesh_device,
+        tt_ccl,
         state_dict,
         state_dict_prefix=first_layer_prefix,
         weight_cache_path=None,

models/tt_transformers/tests/multimodal/mistral_24b/test_vision_attention.py

@@ -8,8 +8,11 @@
 from loguru import logger
 
 import ttnn
-from models.experimental.mistral_24b.tt.vision_attention import TtMistralImageAttention as TtLlamaImageAttention
+from models.tt_transformers.tt.ccl import TT_CCL
 from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.multimodal.mistral_24b.vision_attention import (
+    TtMistralImageAttention as TtLlamaImageAttention,
+)
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 from ttnn import ConcatMeshToTensor
 
@@ -33,6 +36,11 @@
     "batch_size",
     (1,),
 )
+@pytest.mark.parametrize(
+    "device_params",
+    [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "trace_region_size": 30000000, "num_command_queues": 1}],
+    indirect=True,
+)
 def test_vision_attention(mesh_device, seq_len, batch_size):
     logger.info(f"seq_len: {seq_len}, batch_size: {batch_size}")
     dtype = ttnn.bfloat16
@@ -53,8 +61,10 @@ def test_vision_attention(mesh_device, seq_len, batch_size):
     n_heads = model_args.vision_attn_n_heads
     head_dim = hidden_size // n_heads
 
+    tt_ccl = TT_CCL(mesh_device)
     tt_model = TtLlamaImageAttention(
         mesh_device,
+        tt_ccl,
         state_dict,
         state_dict_prefix=first_layer_prefix,
         weight_cache_path=model_args.weight_cache_path(dtype),

models/tt_transformers/tests/multimodal/mistral_24b/test_vision_mlp.py

@@ -9,9 +9,9 @@
 from loguru import logger
 
 import ttnn
-
-from models.experimental.mistral_24b.tt.vision_mlp import MistralTTVisionMLP as MLP
 from models.tt_transformers.tt.model_config import ModelArgs
+
+from models.tt_transformers.tt.multimodal.mistral_24b.vision_mlp import MistralTTVisionMLP as MLP
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 
 

models/tt_transformers/tests/multimodal/mistral_24b/test_vision_model.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_transformers.tt.ccl import TT_CCL
+from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.multimodal.mistral_24b.vision_model import TtMistralVisionTransformer
+from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
+
+
+def get_image_features(vision_tower, projector, input_tensor, image_sizes):
+    """
+    Get image features from the vision tower and projector.
+    """
+    vision_token = vision_tower(input_tensor, image_sizes).last_hidden_state
+    image_features = projector(vision_token.squeeze(0), image_sizes)
+    return image_features
+
+
+@skip_for_grayskull("Requires wormhole_b0 to run")
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "device_params",
+    [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "trace_region_size": 30000000, "num_command_queues": 1}],
+    indirect=True,
+)
+def test_mistral_vision_model(mesh_device, reset_seeds):
+    pcc_required = 0.97
+    dtype = ttnn.bfloat8_b
+
+    model_args = ModelArgs(mesh_device)
+    state_dict = model_args.load_state_dict()
+
+    first_layer_prefix = "vision_tower."
+    partial_state_dict = {
+        k[len(first_layer_prefix) :]: v for k, v in state_dict.items() if k.startswith(first_layer_prefix)
+    }
+
+    ##### Reference model output (Torch) #####
+    reference_model = model_args.reference_vision_model()
+    reference_model.load_state_dict(partial_state_dict)
+
+    mmp_first_layer_prefix = "multi_modal_projector."
+
+    mmp_partial_state_dict = {
+        k[len(mmp_first_layer_prefix) :]: v for k, v in state_dict.items() if (k.startswith(mmp_first_layer_prefix))
+    }
+
+    reference_mmp = model_args.reference_vision_multi_modal()
+    reference_mmp.load_state_dict(mmp_partial_state_dict)
+
+    B, C, H, W = 1, 3, model_args.vision_chunk_size, model_args.vision_chunk_size
+    input_tensor = torch.rand((B, C, H, W), dtype=torch.bfloat16)
+
+    reference_output = get_image_features(reference_model, reference_mmp, input_tensor, image_sizes=[(H, W)])
+
+    # ##### TT Model: TtMistralVisionTransformer #####
+    tt_ccl = TT_CCL(mesh_device=mesh_device)
+    vision_model = TtMistralVisionTransformer(
+        mesh_device=mesh_device,
+        tt_ccl=tt_ccl,
+        state_dict=state_dict,
+        state_dict_prefix=first_layer_prefix,
+        dtype=dtype,
+        model_args=model_args,
+    )
+
+    tt_output = vision_model(input_tensor, image_sizes=[(H, W)])  # [0]
+    tt_output = ttnn.to_torch(tt_output, mesh_composer=ttnn.ConcatMeshToTensor(mesh_device, dim=-1))[
+        :, : tt_output.shape[-1]
+    ]
+
+    non_zero_indices = tt_output.ne(0).nonzero(as_tuple=True)
+    tt_output = tt_output[non_zero_indices]
+    reference_output = reference_output[non_zero_indices]
+
+    passing, pcc_message = comp_pcc(reference_output, tt_output, pcc_required)
+
+    logger.info(comp_allclose(reference_output, tt_output))
+    logger.info(f"PCC: {pcc_message}")
+    assert passing, f"PCC below {pcc_required}. {pcc_message}"

models/tt_transformers/tests/multimodal/mistral_24b/test_vision_rms.py

@@ -5,8 +5,8 @@
 from loguru import logger
 
 import ttnn
-from models.tt_transformers.tt.multimodal.mistral_24b.rmsnorm import RMSNorm
 from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.multimodal.mistral_24b.rmsnorm import RMSNorm
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 
 

models/tt_transformers/tests/multimodal/mistral_24b/test_vision_tower.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.tt_transformers.tt.ccl import TT_CCL
+from models.tt_transformers.tt.model_config import ModelArgs
+from models.tt_transformers.tt.multimodal.mistral_24b.mistral_vision_tower import MistralVisionTower
+from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
+
+
+@skip_for_grayskull("Requires wormhole_b0 to run")
+@pytest.mark.parametrize(
+    "mesh_device",
+    [
+        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
+            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "device_params",
+    [{"fabric_config": ttnn.FabricConfig.FABRIC_1D, "trace_region_size": 30000000, "num_command_queues": 1}],
+    indirect=True,
+)
+def test_mistral_vision_tower(mesh_device, reset_seeds):
+    pcc_required = 0.99
+    dtype = ttnn.bfloat16
+
+    model_args = ModelArgs(mesh_device)
+    state_dict = model_args.load_state_dict()
+
+    first_layer_prefix = "vision_tower."
+    partial_state_dict = {
+        k[len(first_layer_prefix) :]: v for k, v in state_dict.items() if k.startswith(first_layer_prefix)
+    }
+
+    B, C, H, W = 1, 3, model_args.vision_chunk_size, model_args.vision_chunk_size
+    input_tensor = torch.rand((B, C, H, W), dtype=torch.bfloat16)
+
+    ##### Reference model output (Torch) #####
+    reference_model = model_args.reference_vision_model()
+    reference_model.load_state_dict(partial_state_dict)
+    reference_output = reference_model(input_tensor, image_sizes=[(H, W)])
+
+    reference_output = reference_output.last_hidden_state
+    tt_ccl = TT_CCL(mesh_device)
+    ##### TT Model: MistralVisionTower #####
+    vision_model = MistralVisionTower(
+        mesh_device=mesh_device,
+        tt_ccl=tt_ccl,
+        state_dict=state_dict,
+        state_dict_prefix=first_layer_prefix,
+        dtype=dtype,
+        configuration=model_args,
+    )
+
+    tt_output = vision_model(input_tensor, image_sizes=[(H, W)])
+    tt_output = ttnn.to_torch(tt_output, mesh_composer=ttnn.ConcatMeshToTensor(mesh_device, dim=-1))[
+        :, :, :, : tt_output.shape[-1]
+    ]
+    tt_output = tt_output.squeeze(0)
+    passing, pcc_message = comp_pcc(reference_output, tt_output, pcc_required)
+
+    logger.info(comp_allclose(reference_output, tt_output))
+    logger.info(f"PCC: {pcc_message}")
+    assert passing, f"PCC below {pcc_required}. {pcc_message}"