Migrate mistral 24b to tt_transformers

Author: nikileshx

models/tt_transformers/demo/simple_vision_demo.py

@@ -27,7 +27,9 @@
 import ttnn
 from models.demos.utils.llm_demo_utils import create_benchmark_data, verify_perf
 from models.perf.benchmarking_utils import BenchmarkProfiler
+from models.tt_transformers.tt.common import hf_multimodal_encode
 from models.tt_transformers.tt.generator import Generator
+from models.tt_transformers.tt.model_config import CheckpointType
 
 
 def get_batch_sampler(temperature, top_p, tokenizer):
@@ -62,6 +64,7 @@ def create_multimodal_model(
 ):
     from models.tt_transformers.tt.model_config import ModelArgs
     from models.tt_transformers.tt.multimodal.llama_vision_model import CrossAttentionTransformer
+    from models.tt_transformers.tt.multimodal.mistral_24b.mistral_e2e_model import MistralTransformer
 
     tt_model_args = ModelArgs(mesh_device, max_batch_size=max_batch_size)
     assert tt_model_args.is_vision(), "This model is multimodal"
@@ -76,14 +79,25 @@ def create_multimodal_model(
 
     if checkpoint is None:
         checkpoint = tt_model_args.load_state_dict()
-    model = CrossAttentionTransformer(
-        mesh_device,
-        state_dict=checkpoint,
-        weight_cache_path=tt_model_args.weight_cache_path(dtype),
-        dtype=dtype,
-        configuration=tt_model_args,
-        use_paged_kv_cache=use_paged_kv_cache,
-    )
+
+    if tt_model_args.base_model_name == "Mistral-Small-3.1-24B":
+        model = MistralTransformer(
+            mesh_device=mesh_device,
+            state_dict=checkpoint,
+            weight_cache_path=tt_model_args.weight_cache_path(ttnn.bfloat8_b),
+            dtype=ttnn.bfloat8_b,
+            args=tt_model_args,
+            use_paged_kv_cache=use_paged_kv_cache,
+        )
+    else:
+        model = CrossAttentionTransformer(
+            mesh_device,
+            state_dict=checkpoint,
+            weight_cache_path=tt_model_args.weight_cache_path(dtype),
+            dtype=dtype,
+            configuration=tt_model_args,
+            use_paged_kv_cache=use_paged_kv_cache,
+        )
     return tt_model_args, model, checkpoint
 
 
@@ -136,7 +150,7 @@ def prepare_generator_args(
 )
 @pytest.mark.parametrize(
     "test_type,max_seq_len",
-    (("normal", 512),),
+    (("normal", 2048),),
     ids=["normal"],
 )
 @pytest.mark.parametrize(
@@ -182,9 +196,6 @@ def test_multimodal_demo_text(
     profiler = BenchmarkProfiler()
     profiler.start("run")
 
-    ckpt_dir = os.environ["LLAMA_DIR"]
-    tokenizer_path = str(Path(ckpt_dir) / "tokenizer.model")
-
     num_devices = mesh_device.get_num_devices() if isinstance(mesh_device, ttnn.MeshDevice) else 1
     max_batch_size *= data_parallel  # input batch_size is interpreted as size per DP group
 
@@ -195,11 +206,27 @@ def test_multimodal_demo_text(
         max_batch_size=max_batch_size,
         max_seq_len=max_seq_len,
     )
+
+    HF_MODEL = model_args[0].checkpoint_type == CheckpointType.HuggingFace
+
+    if not HF_MODEL:
+        ckpt_dir = os.environ["LLAMA_DIR"]
+        tokenizer_path = str(Path(ckpt_dir) / "tokenizer.model")
+
+        tokenizer = Tokenizer(model_path=tokenizer_path)
+        formatter = ChatFormat(tokenizer)
+    else:
+        from transformers import AutoProcessor
+
+        processor = AutoProcessor.from_pretrained(model_args[0].CKPT_DIR)
+        tokenizer = model_args[0].tokenizer
+
     generator = Generator(model, model_args, mesh_device)
-    tokenizer = Tokenizer(model_path=tokenizer_path)
-    formatter = ChatFormat(tokenizer)
 
-    xattn_caches = [model.setup_cache(model_args[i].max_batch_size) for i, model in enumerate(generator.model)]
+    xattn_caches = [
+        model.setup_cache(model_args[i].max_batch_size) if not HF_MODEL else None
+        for i, model in enumerate(generator.model)
+    ]
 
     # Create random images for trace capture with specific dimensions
     trace_img_560x560 = create_random_image(560, 560)
@@ -264,6 +291,8 @@ def test_multimodal_demo_text(
     _num_prefill_tokens = 0
     _num_decode_tokens = 0
 
+    prompt_encoder = hf_multimodal_encode if HF_MODEL else formatter.encode_dialog_prompt
+
     for iter_num in range(warmup_iters + 1):
         logger.info(f"Iteration {iter_num}")
         current_dialogs = trace_dialogs + dialogs
@@ -273,9 +302,15 @@ def test_multimodal_demo_text(
                 for msg in dialog:
                     print(f"{msg.role.capitalize()}: {msg.content}\n")
             batch_model_input = [
-                formatter.encode_dialog_prompt(dialog, tool_prompt_format=False) for dialog in batch_dialogs
+                prompt_encoder(dialog, processor) if HF_MODEL else prompt_encoder(dialog, tool_prompt_format=False)
+                for dialog in batch_dialogs
             ]
 
+            if HF_MODEL:
+                image_sizes = [model_input.image_sizes for model_input in batch_model_input]
+            else:
+                image_sizes = None
+
             # Do initial prefill
             vision_images = [
                 model_input.vision.images if model_input.vision else None for model_input in batch_model_input
@@ -288,7 +323,7 @@ def test_multimodal_demo_text(
             total_lens = prefill_lens + max_gen_len
 
             # Create padded tokens tensor for batch
-            pad_id = tokenizer.pad_id
+            pad_id = tokenizer.pad_token_id if HF_MODEL else tokenizer.pad_id
             bsz = len(prompt_tokens)
             tokens = torch.full((bsz, max(total_lens)), pad_id, dtype=torch.long)
 
@@ -312,6 +347,7 @@ def test_multimodal_demo_text(
                         xattn_caches,
                         total_lens,
                         prefill_lens,
+                        image_sizes=image_sizes,
                     )
 
             # Get cached prefill time
@@ -323,12 +359,7 @@ def test_multimodal_demo_text(
                     decode_batch_xattn_masks,
                     decode_batch_text_masks,
                 ) = generator.prefill_forward(
-                    vision_images,
-                    vision_mask,
-                    tokens,
-                    xattn_caches,
-                    total_lens,
-                    prefill_lens,
+                    vision_images, vision_mask, tokens, xattn_caches, total_lens, prefill_lens, image_sizes=image_sizes
                 )
 
             prefill_end = time.perf_counter()
@@ -375,12 +406,16 @@ def test_multimodal_demo_text(
                 )  # gen_idx is (num_tokens - 1) to avoid counting compile iter
 
             # Log full text output for each user in batch
-            vision_tokens = [tokenizer.special_tokens["<|image|>"], 128256]
+            if HF_MODEL:
+                # For HF models, get vision tokens from the processor if they exist
+                vision_tokens = []
+            else:
+                vision_tokens = [tokenizer.special_tokens["<|image|>"], 128256]
 
             for user_id in range(max_batch_size):
                 # Remove <|image|> tokens since they break the tokenizer
                 tokens_out = [
-                    t if t not in vision_tokens else tokenizer.pad_id
+                    t if t not in vision_tokens else pad_id
                     for t in tokens[user_id].tolist()[: position_id[user_id] + 2]
                 ]
                 text = tokenizer.decode(tokens_out)

models/experimental/mistral_24b/tests/test_conv2d.py renamed to models/tt_transformers/tests/multimodal/mistral_24b/test_conv2d.py

@@ -12,7 +12,7 @@
 
 import ttnn
 from models.tt_transformers.tt.model_config import ModelArgs
-from models.experimental.mistral_24b.tt.vision_conv2d import TtMistralConv2dPatch
+from models.tt_transformers.tt.multimodal.mistral_24b.vision_conv2d import TtMistralConv2dPatch
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 from ttnn import ConcatMeshToTensor
 

models/experimental/mistral_24b/tests/test_patch_rot_emb.py renamed to models/tt_transformers/tests/multimodal/mistral_24b/test_patch_rot_emb.py

@@ -1,17 +1,17 @@
 # SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 # SPDX-License-Identifier: Apache-2.0
 
-from loguru import logger
+import os
 
-import torch
 import pytest
-import os
-import ttnn
+import torch
+from loguru import logger
 
-from models.experimental.mistral_24b.tt.vision_rope import VisionRotarySetup as RotarySetup
+import ttnn
 
-from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
+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
 
 
 @torch.no_grad()

models/experimental/mistral_24b/tests/test_pixtral_transformer.py renamed to models/tt_transformers/tests/multimodal/mistral_24b/test_pixtral_transformer.py

@@ -8,10 +8,8 @@
 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.experimental.mistral_24b.tt.vision_pixtral_transformer import TtPixtralTransformer
+from models.tt_transformers.tt.model_config import ModelArgs
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 
 
@@ -29,11 +27,6 @@
     ],
     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
 
@@ -58,10 +51,8 @@ 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/experimental/mistral_24b/tests/test_vision_attention.py renamed to models/tt_transformers/tests/multimodal/mistral_24b/test_vision_attention.py

@@ -8,12 +8,9 @@
 from loguru import logger
 
 import ttnn
-from models.tt_transformers.tt.ccl import TT_CCL
+from models.experimental.mistral_24b.tt.vision_attention import TtMistralImageAttention as TtLlamaImageAttention
 from models.tt_transformers.tt.model_config import ModelArgs
 from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
-
-from models.experimental.mistral_24b.tt.vision_attention import TtMistralImageAttention as TtLlamaImageAttention
-
 from ttnn import ConcatMeshToTensor
 
 
@@ -36,11 +33,6 @@
     "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
@@ -61,10 +53,8 @@ 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/experimental/mistral_24b/tests/test_vision_mlp.py renamed to models/tt_transformers/tests/multimodal/mistral_24b/test_vision_mlp.py

@@ -1,19 +1,13 @@
-# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
-
-# SPDX-License-Identifier: Apache-2.0
-
-from loguru import logger
+import os
 
-import torch
 import pytest
-import os
+import torch
+from loguru import logger
 
 import ttnn
-from models.experimental.mistral_24b.tt.rmsnorm import RMSNorm
-
-from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
-
+from models.tt_transformers.tt.multimodal.mistral_24b.rmsnorm import RMSNorm
 from models.tt_transformers.tt.model_config import ModelArgs
+from models.utility_functions import comp_allclose, comp_pcc, skip_for_grayskull
 
 
 @torch.no_grad()

models/experimental/mistral_24b/tests/test_vision_rms.py renamed to models/tt_transformers/tests/multimodal/mistral_24b/test_vision_rms.py

@@ -5,9 +5,11 @@
 import math
 import re
 from enum import Enum
+from types import SimpleNamespace
 from typing import Optional
 
 import torch
+from llama_models.llama3.api.datatypes import ImageMedia
 from loguru import logger
 from pydantic import AliasChoices, BaseModel, Field
 
@@ -688,3 +690,46 @@ def create_tt_model(
     tt_kv_cache = [l.attention.layer_past for l in model.layers] if paged_attention_config else None
 
     return tt_model_args, model, tt_kv_cache, state_dict
+
+
+def hf_multimodal_encode(messages, processor):
+    hf_messages = []
+
+    for msg in messages:
+        hf_content = []
+
+        for item in msg.content:
+            if isinstance(item, ImageMedia):
+                hf_content.append(
+                    {
+                        "type": "image",
+                        "image": item.image,
+                    }
+                )
+            elif isinstance(item, str):
+                hf_content.append(
+                    {
+                        "type": "text",
+                        "text": item,
+                    }
+                )
+
+        hf_messages.append(
+            {
+                "role": msg.role,
+                "content": hf_content,
+            }
+        )
+
+    encoded = processor.apply_chat_template(
+        hf_messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
+    ).to("cpu", dtype=torch.bfloat16)
+
+    return SimpleNamespace(
+        **encoded,
+        tokens=encoded["input_ids"].squeeze(0),
+        vision=SimpleNamespace(
+            images=encoded["pixel_values"],
+            mask=None,
+        ),
+    )