Gemma model files clean up

Author: MohammedTaherMcW

models/tt_transformers/tests/test_attention.py

@@ -164,10 +164,9 @@ def test_attention_inference(
         # 70B attention block typically sees tensors with mean 0 and std 0.03 - 0.05 in layer 1
         pt_attention_input = torch.randn(
             batch_size, seq_len, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
-        ).to(
-            torch.bfloat16
         )  # Qwen2.5 0.5B sees 0.1 to 2.1
-
+        if "gemma" in os.environ.get("HF_MODEL"):
+            pt_attention_input = pt_attention_input.to(torch.bfloat16)
         tt_attention_input = pt_attention_input.clone()
 
         attention_input = model_args.prepare_residual_tensor_decode(

models/tt_transformers/tests/test_attention_prefill.py

@@ -145,7 +145,8 @@ def test_attention_inference(
         )
         * 2
     ) - 1
-    pt_attention_input = pt_attention_input.to(torch.bfloat16)  # Qwen2.5 0.5B sees 0.1 to 2.1
+    if "gemma" in os.environ.get("HF_MODEL"):
+        pt_attention_input = pt_attention_input.to(torch.bfloat16)
     tt_attention_input = pt_attention_input.clone()
     attention_input = model_args.prepare_residual_tensor_prefill(
         tt_attention_input,

models/tt_transformers/tests/test_decoder.py

@@ -168,7 +168,8 @@ def test_decoder_inference(
             )
             * 2
         ) - 1
-        pt_decode_input = pt_decode_input.to(torch.bfloat16)
+        if "gemma" in os.environ.get("HF_MODEL"):
+            pt_decode_input = pt_decode_input.to(torch.bfloat16)
         tt_decode_input = pt_decode_input.clone()
 
         decode_input = model_args.prepare_residual_tensor_decode(

models/tt_transformers/tests/test_decoder_prefill.py

@@ -155,7 +155,8 @@ def test_decoder_inference(
             )
             * 2
         ) - 1
-        pt_decode_input = pt_decode_input.to(torch.bfloat16)
+        if "gemma" in os.environ.get("HF_MODEL"):
+            pt_decode_input = pt_decode_input.to(torch.bfloat16)
         tt_decode_input = pt_decode_input.clone()
         decode_input = model_args.prepare_residual_tensor_prefill(
             tt_decode_input,

models/tt_transformers/tests/test_lm_head.py

@@ -64,7 +64,9 @@ def test_lm_head_inference(seq_len, batch_size, mesh_device, reset_seeds):
         max_columns_per_device=model_args.max_columns_per_device_lm_head,
     )
 
-    torch_input = torch.randn(1, 1, seq_len, model_args.dim).to(torch.bfloat16)
+    torch_input = torch.randn(1, 1, seq_len, model_args.dim)
+    if "gemma" in os.environ.get("HF_MODEL"):
+        torch_input = torch_input.to(torch.bfloat16)
     reference_output = reference_model(torch_input)
     tt_input = ttnn.from_torch(
         torch_input,

models/tt_transformers/tt/model_config.py

@@ -1550,6 +1550,8 @@ def vision_chunk_ntok(self):
         """
         Returns the number of tokens per chunk, accounting for the extra class token
         """
+        if self.is_llama_vision():
+            return (self.vision_chunk_size // self.vision_patch_size) ** 2 + 1
         return (self.image_size // self.vision_patch_size) ** 2 + 1
 
     def _set_model_params(self, checkpoint_dir):
@@ -1683,7 +1685,12 @@ def __repr__(self):
 )"""
 
     # TODO: Rename to is_llama_vision
+    def is_llama_vision(self):
+        return self.vision_chunk_size > 0
+
     def is_vision(self):
+        if self.is_llama_vision():
+            return True
         return self.image_size > 0
 
     def get_state_dict_prefix(self, module_name, layer_num, is_vision=False):