File size: 1,925 Bytes
3e0ec22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | # SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import os
import pytest
from loguru import logger
from models.tt_transformers.tt.common import get_hf_tt_cache_path
# This test will run all the nightly fast dispatch tests for all supported TTT models in CI [N150 / N300 only]
@pytest.mark.parametrize(
"model_weights",
[
"meta-llama/Llama-3.2-1B-Instruct",
"meta-llama/Llama-3.2-3B-Instruct",
"meta-llama/Llama-3.1-8B-Instruct",
"meta-llama/Llama-3.2-11B-Vision-Instruct",
"mistralai/Mistral-7B-Instruct-v0.3",
],
ids=[
"ttt-llama3.2-1B",
"ttt-llama3.2-3B",
"ttt-llama3.1-8B",
"ttt-llama3.2-11B",
"ttt-mistral-7B-v0.3",
],
)
def test_ci_dispatch(model_weights):
logger.info(f"Running fast dispatch tests for {model_weights}")
os.environ["HF_MODEL"] = model_weights
os.environ["TT_CACHE_PATH"] = get_hf_tt_cache_path(model_weights)
# Pass the exit code of pytest to proper keep track of failures during runtime
exit_code = pytest.main(
[
"models/tt_transformers/tests/test_embedding.py",
"models/tt_transformers/tests/test_rms_norm.py",
"models/tt_transformers/tests/test_mlp.py",
"models/tt_transformers/tests/test_attention.py",
"models/tt_transformers/tests/test_attention_prefill.py",
"models/tt_transformers/tests/test_decoder.py",
"models/tt_transformers/tests/test_decoder_prefill.py",
]
+ ["-x"] # Fail if one of the tests fails
+ (["--timeout", "600"] if "mistral" in model_weights.lower() else [])
)
if exit_code == pytest.ExitCode.TESTS_FAILED:
pytest.fail(
f"One or more CI dispatch tests failed for {model_weights}. Please check the log above for more info",
pytrace=False,
)
|