Rebase Gemma-3-4b-it

Author: MohammedTaherMcW

models/tt_transformers/demo/simple_vision_demo.py

@@ -314,6 +314,7 @@ def test_multimodal_demo_text(
             total_lens = prefill_lens + max_gen_len
 
             # Create padded tokens tensor for batch
+            stop_tokens = model_args[0].tokenizer.stop_tokens
             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)
@@ -394,8 +395,14 @@ def test_multimodal_demo_text(
                         profiler.end(f"compile_decode", iteration=batch_idx)
 
                     # Disable checking for eot until I have more robust code for batch > 1
-                    # if text in ["<|eot_id|>", "<|eom_id|>"]:
-                    #     break
+                    if HF_MODEL:
+                        if next_tokens in stop_tokens:
+                            break
+                    else:
+                        # Disable checking for eot until I have more robust code for batch > 1
+                        pass
+                        # if text in ["<|eot_id|>", "<|eom_id|>"]:
+                        #     break
                 _num_decode_tokens += (
                     gen_idx * max_batch_size
                 )  # gen_idx is (num_tokens - 1) to avoid counting compile iter

models/tt_transformers/tests/test_decoder.py

@@ -87,6 +87,19 @@ def test_decoder_inference(
         model_args.rope_theta,
         model_args.rope_scaling,
     )
+
+    if model_args.rope_local_theta is not None:
+        rope_setup_local = RotarySetup(
+            mesh_device,
+            model_args.max_batch_size,
+            model_args.head_dim,
+            model_args.max_seq_len,
+            model_args.rope_local_theta,
+            None,
+        )
+    else:
+        rope_setup_local = None
+
     transformation_mats = rope_setup.get_both_trans_mats()
 
     # Prepare page table for paged attention
@@ -172,12 +185,12 @@ def test_decoder_inference(
 
         # Get cos/sin matrices for the current position of each user
         rot_mats = rope_setup.get_rot_mats(current_pos)
-
+        rot_mats_local = None if rope_setup_local is None else rope_setup_local.get_rot_mats(current_pos)
         # Run TT model
         tt_out = tt_model(
             decode_input,
             current_pos_tensor,
-            rot_mats=rot_mats,
+            rot_mats=[rot_mats, rot_mats_local],
             mode="decode",
             page_table=page_table_tt,
         )

models/tt_transformers/tests/test_decoder_prefill.py

@@ -93,6 +93,16 @@ def test_decoder_inference(
         theta=model_args.rope_theta,
         rope_scaling=model_args.rope_scaling,
     )
+    if model_args.rope_local_theta is not None:
+        rot_mats_local = get_rot_mats(
+            head_dim=model_args.head_dim,
+            device=mesh_device,
+            seq_len=max_seq_len,
+            theta=model_args.rope_local_theta,
+            rope_scaling=None,
+        )
+    else:
+        rot_mats_local = None
     transformation_mat_torch = get_rot_transformation_mat(model_args.head_dim)
     transformation_mats_prefill = ttnn.as_tensor(
         transformation_mat_torch,
@@ -168,7 +178,9 @@ def test_decoder_inference(
         attn_mask_torch = torch.triu(attn_mask, diagonal=1)
         ref_output = reference_model(pt_decode_input, positions[0], freqs_cis_i, mask=attn_mask_torch)
         # Run TT model
-        tt_out = tt_model(decode_input, None, rot_mats, user_id=0, mode="prefill", page_table=page_table_tt)
+        tt_out = tt_model(
+            decode_input, None, [rot_mats, rot_mats_local], user_id=0, mode="prefill", page_table=page_table_tt
+        )
         tt_out = ttnn.to_torch(
             tt_out,
             mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape),

models/tt_transformers/tests/test_embedding.py

@@ -42,7 +42,7 @@ def test_embedding(max_seq_len, batch_size, mesh_device, reset_seeds, ensure_gc)
     tokenizer = model_args.tokenizer
 
     reference_emb = model_args.reference_embedding()
-    if model_args.is_vision():
+    if model_args.is_vision() and not model_args.base_model_name.startswith("gemma-3"):
         layer_name = "text_model.tok_embeddings.weight"
     else:
         layer_name = "tok_embeddings.weight"
@@ -68,7 +68,8 @@ def test_embedding(max_seq_len, batch_size, mesh_device, reset_seeds, ensure_gc)
         dtype=ttnn.uint32,
         layout=ttnn.ROW_MAJOR_LAYOUT,
     )
-    tt_output = tt_emb(tt_input)
+    embed_scale = model_args.embed_scale
+    tt_output = tt_emb(tt_input, embed_scale)
     tt_output_torch = ttnn.to_torch(
         tt_output,
         mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(0, -1), mesh_shape=model_args.cluster_shape),

models/tt_transformers/tt/common.py

@@ -4,7 +4,9 @@
 
 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

models/tt_transformers/tt/generator.py

@@ -59,7 +59,6 @@ def __init__(self, model, model_args, mesh_device, tokenizer=None, formatter=Non
     def prefill_forward_text(
         self, tokens: torch.Tensor, page_table=None, kv_cache=None, prompt_lens=None, empty_slots=None, **kwargs
     ):
-        print("prefill generator ", kwargs["processed_inputs"])
         if page_table is not None:
             assert isinstance(page_table, torch.Tensor), "page_table mush be torch.Tensor"
 

models/tt_transformers/tt/model.py

@@ -61,14 +61,12 @@ def __init__(
 
         if args.rope_local_theta is not None:
             self.rope_setup_local = ActualRopeSetupClass(
-                mesh_device,
-                args.max_batch_size,
-                args.head_dim,
-                args.max_seq_len,
-                args.rope_local_theta,
-                args.rope_scaling_factor,
-                args.orig_context_len,
-                "default",
+                device=mesh_device,
+                batch_size=args.max_batch_size,
+                head_dim=args.head_dim,
+                max_seq_len=args.max_seq_len,
+                rope_theta=args.rope_local_theta,
+                rope_scaling=None,
             )
         else:
             self.rope_setup_local = None

models/tt_transformers/tt/model_config.py

@@ -1746,20 +1746,6 @@ def merge_vision_config(base_config):
                         self._set_vision_params(merged_vision_config)
             else:
                 self._set_params_from_dict(self.hf_config, is_hf=True)
-
-            if "text_config" in config or "vision_config" in config:
-                merged_text_config = merge_text_config(config)
-                self._set_params_from_dict(merged_text_config, is_hf=True)
-
-                if "gemma-3-4b-it" in self.base_model_name:
-                    self._set_vision_params(config["vision_config"])
-                else:
-                    if "vision_config" in config:
-                        merged_vision_config = merge_vision_config(config)
-                        self._set_vision_params(merged_vision_config)
-            else:
-                self._set_params_from_dict(config, is_hf=True)
-
         else:
             config_file = os.path.join(checkpoint_dir, "config.json")
             assert os.path.exists(config_file), f"config.json file not found at {config_file}"
@@ -2343,9 +2329,6 @@ def reference_transformer(self, wrap=True, load_checkpoint=False):
             # Special case Qwen2.5-VL models until they are fully integrated into a HF release
             if "Qwen/Qwen2.5-VL" in self.model_name:
                 from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig as AutoConfig
-                from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
-                    Qwen2_5_VLForConditionalGeneration as AutoModelForCausalLM,
-                )
             else:
                 from transformers import AutoConfig, AutoModel
 

models/tt_transformers/tt/multimodal/gemma/gemma_image_attention.py

@@ -273,7 +273,7 @@ def pad_head_dim_bias(bias):
                 memory_config=ttnn.DRAM_MEMORY_CONFIG,
                 dtype=self.dtype,
                 layout=ttnn.TILE_LAYOUT,
-                # cache_file_name=cache_name("bo_sharded"),
+                cache_file_name=cache_name("bo_sharded"),
             )
         else:
             self.bo = None

models/tt_transformers/tt/rope.py

@@ -335,7 +335,6 @@ def __init__(
         self.batch_size = batch_size
         self.head_dim = head_dim
         self.device = device
-        self.rope_type = rope_type
         self.is_mesh_device = isinstance(device, ttnn._ttnn.multi_device.MeshDevice)
         self.num_devices = device.get_num_devices() if self.is_mesh_device else 1
         if self.num_devices == 32: