#7502: Move test_bert_batch_dram.py totally into models 2 script beca…

…use we run on both archs with some skips
tenstorrent · Apr 17, 2024 · 309484e · 309484e
1 parent eafbbc0
commit 309484e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/models/demos/metal_BERT_large_11/tests/test_bert_batch_dram.py b/models/demos/metal_BERT_large_11/tests/test_bert_batch_dram.py
@@ -21,6 +21,7 @@
     disable_persistent_kernel_cache,
     profiler,
     is_e75,
+    skip_for_wormhole_b0,
 )
 
 
@@ -282,6 +283,9 @@ def test_bert_batch_dram(
     if is_e75(device):
         pytest.skip(f"Bert large 11 is not supported on E75")
 
+    if device.arch() == tt_lib.device.Arch.WORMHOLE_B0 and (batch != 7 or model_config_str != "BFLOAT8_B-SHARDED"):
+        pytest.skip("Only batch_7-BFLOAT8_B-SHARDED supported for WH B0")
+
     model_config = get_model_config(batch, device.compute_with_storage_grid_size(), model_config_str)
     tt_cache_path = get_tt_cache_path(model_version)
 
@@ -365,6 +369,9 @@ def test_bert_batch_dram_with_program_cache(
     if is_e75(device):
         pytest.skip(f"Bert large 11 is not supported on E75")
 
+    if device.arch() == tt_lib.device.Arch.WORMHOLE_B0 and (batch != 7 or model_config_str != "BFLOAT8_B-SHARDED"):
+        pytest.skip("Only batch_7-BFLOAT8_B-SHARDED supported for WH B0")
+
     model_config = get_model_config(batch, device.compute_with_storage_grid_size(), model_config_str)
     tt_cache_path = get_tt_cache_path(model_version)
 

diff --git a/tests/scripts/nightly/run_models_1.sh b/tests/scripts/nightly/run_models_1.sh
@@ -35,7 +35,6 @@ env pytest models/experimental/stable_diffusion/tests/test_unet_mid_block.py -k
 env pytest models/experimental/stable_diffusion/tests/test_upblock_2d.py -k test_run_upblock_real_input_inference
 env pytest models/experimental/stable_diffusion/tests -k test_unbatched_stable_diffusion #
 
-env pytest models/demos/metal_BERT_large_11/tests/test_bert_batch_dram.py
 env pytest models/demos/metal_BERT_large_11/tests/test_demo.py::test_demo -k batch_12
 env pytest models/demos/metal_BERT_large_11/tests/test_demo.py::test_demo_squadv2 -k batch_12
 

diff --git a/tests/scripts/nightly/run_models_2.sh b/tests/scripts/nightly/run_models_2.sh
@@ -24,7 +24,7 @@ env pytest models/experimental/bloom/tests/test_bloom_gelu_forward.py
 env pytest models/experimental/bloom/tests/test_bloom_merge_heads.py
 env pytest models/experimental/bloom/tests/test_bloom_mlp.py
 
-env pytest models/demos/metal_BERT_large_11/tests/test_bert_batch_dram.py -k batch_7-BFLOAT8_B-SHARDED
+env pytest models/demos/metal_BERT_large_11/tests/test_bert_batch_dram.py
 env pytest models/demos/metal_BERT_large_11/tests/test_demo.py
 
 env pytest models/experimental/synthetic_gradients/tests/