#5337: Update Mistral-7B CI flags

tenstorrent · Jun 5, 2024 · 692bdf4 · 692bdf4
1 parent e816db1
commit 692bdf4
Show file tree

Hide file tree

Showing 8 changed files with 10 additions and 8 deletions.
diff --git a/models/demos/wormhole/mistral7b/demo/demo.py b/models/demos/wormhole/mistral7b/demo/demo.py
@@ -13,7 +13,6 @@
     os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
-    os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
     os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn
@@ -111,7 +110,7 @@ def run_mistral_demo(user_input, batch_size, device, instruct_mode):
     model_args = TtModelArgs(device, instruct=instruct_mode)
     tokenizer = Tokenizer(model_args.tokenizer_path)
 
-    model_args.n_layers = 1
+    model_args.n_layers = 32
 
     logger.info("Loading weights...")
     state_dict = torch.load(model_args.consolidated_weights_path)

diff --git a/models/demos/wormhole/mistral7b/tests/test_mistral_attention.py b/models/demos/wormhole/mistral7b/tests/test_mistral_attention.py
@@ -11,7 +11,6 @@
     os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
-    os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
     os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn

diff --git a/models/demos/wormhole/mistral7b/tests/test_mistral_decoder.py b/models/demos/wormhole/mistral7b/tests/test_mistral_decoder.py
@@ -11,7 +11,6 @@
     os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
-    os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
     os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn

diff --git a/models/demos/wormhole/mistral7b/tests/test_mistral_embedding.py b/models/demos/wormhole/mistral7b/tests/test_mistral_embedding.py
@@ -11,7 +11,6 @@
     os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
-    os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
     os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn

diff --git a/models/demos/wormhole/mistral7b/tests/test_mistral_mlp.py b/models/demos/wormhole/mistral7b/tests/test_mistral_mlp.py
@@ -12,7 +12,6 @@
     os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
-    os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
     os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn

diff --git a/models/demos/wormhole/mistral7b/tests/test_mistral_model.py b/models/demos/wormhole/mistral7b/tests/test_mistral_model.py
@@ -4,6 +4,15 @@
 import torch
 import pytest
 from loguru import logger
+import os
+
+# Set Mistral flags for CI, if CI environment is setup
+if os.getenv("CI") == "true":
+    os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
+    os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
+    os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
+    os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
+
 import ttnn
 from models.demos.wormhole.mistral7b.tt.mistral_common import (
     precompute_freqs,

diff --git a/models/demos/wormhole/mistral7b/tests/test_mistral_perf.py b/models/demos/wormhole/mistral7b/tests/test_mistral_perf.py
@@ -11,7 +11,6 @@
     os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
-    os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
     os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn

diff --git a/models/demos/wormhole/mistral7b/tests/test_mistral_rms_norm.py b/models/demos/wormhole/mistral7b/tests/test_mistral_rms_norm.py
@@ -11,7 +11,6 @@
     os.environ["MISTRAL_CKPT_DIR"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_TOKENIZER_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
     os.environ["MISTRAL_CACHE_PATH"] = "/mnt/MLPerf/ttnn/models/demos/mistral7b/"
-    os.environ["TT_METAL_ASYNC_DEVICE_QUEUE"] = "1"
     os.environ["WH_ARCH_YAML"] = "wormhole_b0_80_arch_eth_dispatch.yaml"
 
 import ttnn