Add BLOOM tests

tests/jax/models/bloom/bloom_1b1/test_1b1.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable
+
+import pytest
+from infra import ModelTester, RunMode
+from utils import compile_fail, record_model_test_properties
+
+from ..tester import BloomTester
+
+MODEL_PATH = "bigscience/bloom-1b1"
+MODEL_NAME = "bloom-1.1b"
+
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> BloomTester:
+    return BloomTester(MODEL_PATH)
+
+
+@pytest.fixture
+def training_tester() -> BloomTester:
+    return BloomTester(ModelTester, run_mode=RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+# This is an interesting one.
+# The error message seems to happen before the compile even begins
+# And then then compile segfaults with no useful information
+# It is highly likely that both are caused by the same root cause
+@pytest.mark.skip(reason=compile_fail("Unsupported data type"))  # segfault
+def test_bloom_1b1_inference(
+    inference_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    inference_tester.test()
+
+
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_bloom_1b1_training(
+    training_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    training_tester.test()

tests/jax/models/bloom/bloom_1b7/test_1b7.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable
+
+import pytest
+from infra import ModelTester, RunMode
+from utils import compile_fail, record_model_test_properties
+
+from ..tester import BloomTester
+
+MODEL_PATH = "bigscience/bloom-1b7"
+MODEL_NAME = "bloom-1.7b"
+
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> BloomTester:
+    return BloomTester(MODEL_PATH)
+
+
+@pytest.fixture
+def training_tester() -> BloomTester:
+    return BloomTester(ModelTester, run_mode=RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+
+@pytest.mark.skip(reason=compile_fail("Unsupported data type"))  # segfault
+def test_bloom_1b7_inference(
+    inference_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    inference_tester.test()
+
+
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_bloom_1b7_training(
+    training_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    training_tester.test()

tests/jax/models/bloom/bloom_3b/test_3b.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable
+
+import pytest
+from infra import ModelTester, RunMode
+from utils import compile_fail, record_model_test_properties
+
+from ..tester import BloomTester
+
+MODEL_PATH = "bigscience/bloom-3b"
+MODEL_NAME = "bloom-3b"
+
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> BloomTester:
+    return BloomTester(MODEL_PATH)
+
+
+@pytest.fixture
+def training_tester() -> BloomTester:
+    return BloomTester(ModelTester, run_mode=RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+
+@pytest.mark.skip(reason=compile_fail("Unsupported data type"))  # segfault
+def test_bloom_3b_inference(
+    inference_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    inference_tester.test()
+
+
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_bloom_3b_training(
+    training_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    training_tester.test()

tests/jax/models/bloom/bloom_560m/test_560m.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable
+
+import pytest
+from infra import ModelTester, RunMode
+from utils import compile_fail, record_model_test_properties
+
+from ..tester import BloomTester
+
+MODEL_PATH = "bigscience/bloom-560m"
+MODEL_NAME = "bloom-560m"
+
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> BloomTester:
+    return BloomTester(MODEL_PATH)
+
+
+@pytest.fixture
+def training_tester() -> BloomTester:
+    return BloomTester(ModelTester, run_mode=RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+
+@pytest.mark.skip(reason=compile_fail("Unsupported data type"))  # segfault
+def test_bloom_560m_inference(
+    inference_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    inference_tester.test()
+
+
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_bloom_560m_training(
+    training_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    training_tester.test()

tests/jax/models/bloom/bloom_7b/test_7b.py

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable
+
+import pytest
+from infra import ModelTester, RunMode
+from utils import compile_fail, record_model_test_properties
+
+from ..tester import BloomTester
+
+MODEL_PATH = "bigscience/bloom-7b1"
+MODEL_NAME = "bloom-7b"
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> BloomTester:
+    return BloomTester(MODEL_PATH)
+
+
+@pytest.fixture
+def training_tester() -> BloomTester:
+    return BloomTester(ModelTester, run_mode=RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+
+@pytest.mark.skip(reason=compile_fail("Unsupported data type"))  # segfault
+def test_bloom_7b_inference(
+    inference_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    inference_tester.test()
+
+
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_bloom_7b_training(
+    training_tester: BloomTester,
+    record_tt_xla_property: Callable,
+):
+    record_model_test_properties(record_tt_xla_property, MODEL_NAME)
+
+    training_tester.test()

tests/jax/models/bloom/tester.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Dict, Sequence
+
+import jax
+from flax import linen as nn
+from infra import ComparisonConfig, ModelTester, RunMode
+from transformers import AutoTokenizer, FlaxBloomForCausalLM
+
+
+class BloomTester(ModelTester):
+    """Tester for Bloom models."""
+
+    def __init__(
+        self,
+        model_name: str,
+        comparison_config: ComparisonConfig = ComparisonConfig(),
+        run_mode: RunMode = RunMode.INFERENCE,
+    ) -> None:
+        self._model_name = model_name
+        super().__init__(comparison_config, run_mode)
+
+    # @override
+    def _get_model(self) -> nn.Module:
+        return FlaxBloomForCausalLM.from_pretrained(self._model_name, from_pt=True)
+
+    # @override
+    def _get_input_activations(self) -> Sequence[jax.Array]:
+        tokenizer = AutoTokenizer.from_pretrained(self._model_name)
+        inputs = tokenizer("Hello", return_tensors="np")
+        return inputs["input_ids"]
+
+    # @override
+    def _get_forward_method_kwargs(self) -> Dict[str, jax.Array]:
+        assert hasattr(self._model, "params")
+        return {
+            "params": self._model.params,
+            "input_ids": self._get_input_activations(),
+        }