#7502: models 1 - all clean GS-only model tests with skip_for_wormhole

models/demos/metal_BERT_large_11/tests/test_demo.py

@@ -6,9 +6,10 @@
 from models.demos.metal_BERT_large_11.demo.demo import test_demo_squadv2 as demo_squadv2
 import pytest
 from loguru import logger
-from models.utility_functions import is_e75, skip_for_wormhole_b0
+from models.utility_functions import is_e75, skip_for_wormhole_b0, skip_for_grayskull
 
 
+@skip_for_grayskull()
 @pytest.mark.parametrize("batch", (7,), ids=["batch_7"])
 @pytest.mark.parametrize(
     "input_path",
@@ -78,6 +79,7 @@ def test_demo_batch_12(batch, input_path, model_location_generator, device, use_
         assert expected_answers[i] == answers[i]
 
 
+@skip_for_grayskull()
 @pytest.mark.parametrize(
     "batch, exact, f1",
     (

models/experimental/bloom/tests/test_bloom_block.py

@@ -13,6 +13,7 @@
 from loguru import logger
 import models.experimental.bloom.bloom_utils as bloom_utils
 import models.experimental.bloom.tt.bloom_block as bloom_block
+from models.utility_functions import skip_for_wormhole_b0
 
 
 def run_bloom_block_test(device):
@@ -61,5 +62,6 @@ def run_bloom_block_test(device):
     assert do_all_blocks_pass
 
 
+@skip_for_wormhole_b0()
 def test_bloom_block(device):
     run_bloom_block_test(device)

models/experimental/bloom/tests/test_bloom_causal_lm.py

@@ -13,12 +13,12 @@
 
 from loguru import logger
 import models.experimental.bloom.tt.bloom_causal_lm as bloom_causal_lm
+from models.utility_functions import skip_for_wormhole_b0
 
 
+@skip_for_wormhole_b0()
 def test_bloom_causal_lm(device):
-    hugging_bloom_reference_model = BloomForCausalLM.from_pretrained(
-        "bigscience/bloom-560m", torchscript=False
-    )
+    hugging_bloom_reference_model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m", torchscript=False)
     hugging_bloom_reference_model.eval()
 
     config = hugging_bloom_reference_model.config

models/experimental/bloom/tests/test_bloom_model.py

@@ -15,12 +15,11 @@
 
 import models.experimental.bloom.bloom_utils as bloom_utils
 import models.experimental.bloom.tt.bloom_model as bloom_model
+from models.utility_functions import skip_for_wormhole_b0
 
 
 def run_bloom_model_test(device):
-    hugging_bloom_reference_model = BloomForCausalLM.from_pretrained(
-        "bigscience/bloom-560m", torchscript=False
-    )
+    hugging_bloom_reference_model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m", torchscript=False)
     hugging_bloom_reference_model.eval()
 
     config = hugging_bloom_reference_model.config
@@ -56,5 +55,6 @@ def run_bloom_model_test(device):
     assert does_pass
 
 
+@skip_for_wormhole_b0()
 def test_bloom_model(device):
     run_bloom_model_test(device)

models/experimental/mnist/tests/test_mnist.py

@@ -13,10 +13,12 @@
     torch2tt_tensor,
     tt2torch_tensor,
     comp_pcc,
+    skip_for_wormhole_b0,
 )
 from models.experimental.mnist.tt.mnist_model import mnist_model
 
 
+@skip_for_wormhole_b0()
 def test_mnist_inference(device, model_location_generator):
     # Data preprocessing/loading
     transform = transforms.Compose([transforms.ToTensor()])

models/experimental/nanogpt/tests/test_nanogpt_mlp.py

@@ -20,6 +20,7 @@
     torch_to_tt_tensor_rm,
     comp_allclose,
     comp_pcc,
+    skip_for_wormhole_b0,
 )
 
 
@@ -31,6 +32,7 @@
     "pcc",
     ((0.99,),),
 )
+@skip_for_wormhole_b0()
 def test_nanogpt_mlp(device, pcc, dtype, reset_seeds):
     model_hf = GPT2LMHeadModel.from_pretrained("gpt2")
     config = model_hf.config

models/experimental/roberta/tests/test_roberta_encoder.py

@@ -8,16 +8,14 @@
 from transformers import RobertaModel
 
 import tt_lib
+import pytest
 
 from models.experimental.roberta.tt.roberta_encoder import TtRobertaEncoder
-from models.utility_functions import (
-
-    tt2torch_tensor,
-    comp_allclose,
-    comp_pcc,
-)
+from models.utility_functions import tt2torch_tensor, comp_allclose, comp_pcc, skip_for_wormhole_b0
 from models.experimental.roberta.roberta_common import torch2tt_tensor
 
+
+@skip_for_wormhole_b0()
 def test_roberta_encoder_inference(device):
     torch.manual_seed(1234)
     base_address = f"encoder"
@@ -51,9 +49,7 @@ def test_roberta_encoder_inference(device):
     tt_output_torch = tt2torch_tensor(tt_output.last_hidden_state)
     tt_output_torch = tt_output_torch.squeeze(0)
 
-    does_pass, pcc_message = comp_pcc(
-        torch_output.last_hidden_state, tt_output_torch, 0.98
-    )
+    does_pass, pcc_message = comp_pcc(torch_output.last_hidden_state, tt_output_torch, 0.98)
 
     logger.info(comp_allclose(torch_output.last_hidden_state, tt_output_torch))
     logger.info(pcc_message)

models/experimental/roberta/tests/test_roberta_for_masked_lm.py

@@ -8,16 +8,19 @@
 from transformers import AutoTokenizer, RobertaForMaskedLM
 
 import tt_lib
+import pytest
 
 from models.experimental.roberta.tt.roberta_for_masked_lm import TtRobertaForMaskedLM
 from models.utility_functions import (
     tt2torch_tensor,
     comp_allclose,
     comp_pcc,
+    skip_for_wormhole_b0,
 )
 from models.experimental.roberta.roberta_common import torch2tt_tensor
 
 
+@skip_for_wormhole_b0()
 def test_roberta_masked_lm_inference(device):
     torch.manual_seed(1234)
     base_address = f""
@@ -47,9 +50,7 @@ def test_roberta_masked_lm_inference(device):
 
         torch_output = torch_output.logits
         # retrieve index of <mask>
-        mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(
-            as_tuple=True
-        )[0]
+        mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
 
         predicted_token_id = torch_output[0, mask_token_index].argmax(axis=-1)
         decoded_token = tokenizer.decode(predicted_token_id)
@@ -71,9 +72,7 @@ def test_roberta_masked_lm_inference(device):
         tt_output_torch = tt_output_torch.squeeze(0)
 
         # retrieve index of <mask>
-        mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(
-            as_tuple=True
-        )[0]
+        mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
         # Get predicted token
         predicted_token_id = tt_output_torch[0, mask_token_index].argmax(axis=-1)
         decoded_token = tokenizer.decode(predicted_token_id)

models/experimental/roberta/tests/test_roberta_for_multiple_choice.py

@@ -15,10 +15,12 @@
     tt2torch_tensor,
     comp_allclose,
     comp_pcc,
+    skip_for_wormhole_b0,
 )
 from models.experimental.roberta.roberta_common import torch2tt_tensor
 
 
+@skip_for_wormhole_b0()
 @pytest.mark.skip(reason="Mismatch happening on GS, issue #5943")
 def test_roberta_for_multiple_choice(device):
     """

models/experimental/roberta/tests/test_roberta_for_qa.py

@@ -8,23 +8,21 @@
 from transformers import AutoTokenizer, RobertaForQuestionAnswering
 
 import tt_lib
+import pytest
 
 from models.experimental.roberta.tt.roberta_for_question_answering import TtRobertaForQuestionAnswering
-from models.utility_functions import (
-    comp_allclose,
-    comp_pcc,
-)
+from models.utility_functions import comp_allclose, comp_pcc, skip_for_wormhole_b0
 from models.experimental.roberta.roberta_common import torch2tt_tensor
 
+
+@skip_for_wormhole_b0()
 def test_roberta_qa_inference(device):
     torch.manual_seed(1234)
 
     base_address = f""
 
     tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
-    torch_model = RobertaForQuestionAnswering.from_pretrained(
-        "deepset/roberta-base-squad2"
-    )
+    torch_model = RobertaForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
     torch_model.eval()
 
     tt_model = TtRobertaForQuestionAnswering(
@@ -49,12 +47,8 @@ def test_roberta_qa_inference(device):
         torch_answer_start_index = torch_output.start_logits.argmax()
         torch_answer_end_index = torch_output.end_logits.argmax()
 
-        torch_predict_answer_tokens = inputs.input_ids[
-            0, torch_answer_start_index : torch_answer_end_index + 1
-        ]
-        torch_answer = tokenizer.decode(
-            torch_predict_answer_tokens, skip_special_tokens=True
-        )
+        torch_predict_answer_tokens = inputs.input_ids[0, torch_answer_start_index : torch_answer_end_index + 1]
+        torch_answer = tokenizer.decode(torch_predict_answer_tokens, skip_special_tokens=True)
         logger.info("Torch answered")
         logger.info(torch_answer)
 
@@ -68,24 +62,18 @@ def test_roberta_qa_inference(device):
         tt_answer_start_index = tt_output.start_logits.argmax()
         tt_answer_end_index = tt_output.end_logits.argmax()
 
-        tt_predict_answer_tokens = inputs.input_ids[
-            0, tt_answer_start_index : tt_answer_end_index + 1
-        ]
+        tt_predict_answer_tokens = inputs.input_ids[0, tt_answer_start_index : tt_answer_end_index + 1]
         tt_answer = tokenizer.decode(tt_predict_answer_tokens, skip_special_tokens=True)
         logger.info("TT answered")
         logger.info(tt_answer)
 
         # Compare outputs
-        does_pass_1, pcc_message = comp_pcc(
-            torch_output.start_logits, tt_output.start_logits, 0.98
-        )
+        does_pass_1, pcc_message = comp_pcc(torch_output.start_logits, tt_output.start_logits, 0.98)
 
         logger.info(comp_allclose(torch_output.start_logits, tt_output.start_logits))
         logger.info(pcc_message)
 
-        does_pass_2, pcc_message = comp_pcc(
-            torch_output.end_logits, tt_output.end_logits, 0.98
-        )
+        does_pass_2, pcc_message = comp_pcc(torch_output.end_logits, tt_output.end_logits, 0.98)
 
         logger.info(comp_allclose(torch_output.end_logits, tt_output.end_logits))
         logger.info(pcc_message)

models/experimental/roberta/tests/test_roberta_for_sequence_classification.py

@@ -8,26 +8,26 @@
 from transformers import AutoTokenizer, RobertaForSequenceClassification
 
 import tt_lib
+import pytest
 
 from models.experimental.roberta.tt.roberta_for_sequence_classification import TtRobertaForSequenceClassification
 from models.utility_functions import (
     tt2torch_tensor,
     comp_allclose,
     comp_pcc,
+    skip_for_wormhole_b0,
 )
 from models.experimental.roberta.roberta_common import torch2tt_tensor
 
+
+@skip_for_wormhole_b0()
 def test_roberta_for_sequence_classification(device):
     torch.manual_seed(1234)
     base_address = ""
 
     with torch.no_grad():
-        tokenizer = AutoTokenizer.from_pretrained(
-            "cardiffnlp/twitter-roberta-base-emotion"
-        )
-        model = RobertaForSequenceClassification.from_pretrained(
-            "cardiffnlp/twitter-roberta-base-emotion"
-        )
+        tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-emotion")
+        model = RobertaForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-emotion")
         model.eval()
 
         # Tt roberta

models/experimental/roberta/tests/test_roberta_layer.py

@@ -8,15 +8,19 @@
 from transformers import RobertaModel
 
 import tt_lib
+import pytest
 
 from models.experimental.roberta.tt.roberta_layer import TtRobertaLayer
 from models.utility_functions import (
     tt2torch_tensor,
     comp_allclose,
     comp_pcc,
+    skip_for_wormhole_b0,
 )
 from models.experimental.roberta.roberta_common import torch2tt_tensor
 
+
+@skip_for_wormhole_b0()
 def test_roberta_layer_inference(device):
     torch.manual_seed(1234)
 

models/experimental/roberta/tests/test_roberta_model.py

@@ -8,15 +8,19 @@
 from transformers import AutoTokenizer, RobertaModel
 
 import tt_lib
+import pytest
 
 from models.experimental.roberta.tt.roberta_model import TtRobertaModel
 from models.utility_functions import (
     tt2torch_tensor,
     comp_allclose,
     comp_pcc,
+    skip_for_wormhole_b0,
 )
 from models.experimental.roberta.roberta_common import torch2tt_tensor
 
+
+@skip_for_wormhole_b0()
 def test_roberta_model_inference(device):
     torch.manual_seed(1234)
 

models/experimental/stable_diffusion/tests/test_downblock_2d.py

@@ -11,15 +11,19 @@
 
 
 import tt_lib as ttl
+import pytest
+
 from models.utility_functions import (
     torch_to_tt_tensor,
     tt_to_torch_tensor,
     torch_to_tt_tensor_rm,
+    skip_for_wormhole_b0,
 )
 from models.utility_functions import comp_pcc, comp_allclose_and_pcc
 from models.experimental.stable_diffusion.tt.downblock_2d import TtDownBlock2D
 
 
+@skip_for_wormhole_b0()
 def test_run_downblock_real_input_inference(device, model_location_generator):
     # setup pytorch model
     pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
@@ -62,6 +66,7 @@ def test_run_downblock_real_input_inference(device, model_location_generator):
     logger.info(f"PASSED {passing[1]}")
 
 
+@pytest.mark.skip("Test not run")
 def test_run_downblock_inference(device):
     # setup pytorch model
     pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)

models/experimental/stable_diffusion/tests/test_upblock_2d.py

@@ -12,12 +12,14 @@
     torch_to_tt_tensor,
     tt_to_torch_tensor,
     torch_to_tt_tensor_rm,
+    skip_for_wormhole_b0,
 )
 from models.utility_functions import comp_pcc, comp_allclose_and_pcc
 from models.experimental.stable_diffusion.tt.upblock_2d import TtUpBlock2D
 import pytest
 
 
+@skip_for_wormhole_b0()
 def test_run_upblock_real_input_inference(device, model_location_generator):
     # setup pytorch model
     pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
@@ -65,6 +67,7 @@ def test_run_upblock_real_input_inference(device, model_location_generator):
     logger.info(f"PASSED {passing[1]}")
 
 
+@pytest.mark.skip(reason="Test not run")
 def test_run_upblock_inference(device):
     # setup pytorch model
     pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)