unslothai
diff --git a/‎tests/__init__.py‎ b/‎tests/__init__.py‎
diff --git a/‎tests/test_model_registry.py‎
Lines changed: 91 additions & 0 deletions b/‎tests/test_model_registry.py‎
Lines changed: 91 additions & 0 deletions
diff --git a/‎unsloth/_auto_install.py‎
Lines changed: 2 additions & 1 deletion b/‎unsloth/_auto_install.py‎
Lines changed: 2 additions & 1 deletion
@@ -0,0 +1,91 @@
+"""
+
+Test model registration methods
+Checks that model registration methods work for respective models as well as all models
+The check is performed
+- by registering the models
+- checking that the instantiated models can be found on huggingface hub by querying for the model id
+
+"""
+
+from dataclasses import dataclass
+
+import pytest
+from huggingface_hub import ModelInfo as HfModelInfo
+
+from unsloth.registry import register_models, search_models
+from unsloth.registry._deepseek import register_deepseek_models
+from unsloth.registry._gemma import register_gemma_models
+from unsloth.registry._llama import register_llama_models
+from unsloth.registry._mistral import register_mistral_models
+from unsloth.registry._phi import register_phi_models
+from unsloth.registry._qwen import register_qwen_models
+from unsloth.registry.registry import MODEL_REGISTRY, QUANT_TAG_MAP, QuantType
+from unsloth.utils.hf_hub import get_model_info
+
+MODEL_NAMES = [
+    "llama",
+    "qwen",
+    "mistral",
+    "phi",
+    "gemma",
+    "deepseek",
+]
+MODEL_REGISTRATION_METHODS = [
+    register_llama_models,
+    register_qwen_models,
+    register_mistral_models,
+    register_phi_models,
+    register_gemma_models,
+    register_deepseek_models,
+]
+
+
+@dataclass
+class ModelTestParam:
+    name: str
+    register_models: callable
+
+
+def _test_model_uploaded(model_ids: list[str]):
+    missing_models = []
+    for _id in model_ids:
+        model_info: HfModelInfo = get_model_info(_id)
+        if not model_info:
+            missing_models.append(_id)
+
+    return missing_models
+
+
+TestParams = [
+    ModelTestParam(name, models)
+    for name, models in zip(MODEL_NAMES, MODEL_REGISTRATION_METHODS)
+]
+
+
+# Test that model registration methods register respective models
+@pytest.mark.parametrize("model_test_param", TestParams, ids=lambda param: param.name)
+def test_model_registration(model_test_param: ModelTestParam):
+    MODEL_REGISTRY.clear()
+    registration_method = model_test_param.register_models
+    registration_method()
+    registered_models = MODEL_REGISTRY.keys()
+    missing_models = _test_model_uploaded(registered_models)
+    assert not missing_models, (
+        f"{model_test_param.name} missing following models: {missing_models}"
+    )
+
+
+def test_all_model_registration():
+    register_models()
+    registered_models = MODEL_REGISTRY.keys()
+    missing_models = _test_model_uploaded(registered_models)
+    assert not missing_models, f"Missing following models: {missing_models}"
+
+def test_quant_type():
+    # Test that the quant_type is correctly set for model paths
+    # NOTE: for models registered under org="unsloth" with QuantType.NONE aliases QuantType.UNSLOTH
+    dynamic_quant_models = search_models(quant_types=[QuantType.UNSLOTH])
+    assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models)
+    quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH]
+    assert all(quant_tag in m.model_path for m in dynamic_quant_models)
@@ -18,7 +18,7 @@
 v = V(torch.__version__)
 cuda = str(torch.version.cuda)
 is_ampere = torch.cuda.get_device_capability()[0] >= 8
-if cuda != "12.1" and cuda != "11.8" and cuda != "12.4" and cuda != "12.6": raise RuntimeError(f"CUDA = {cuda} not supported!")
+if cuda != "12.1" and cuda != "11.8" and cuda != "12.4" and cuda != "12.6" and cuda != "12.8": raise RuntimeError(f"CUDA = {cuda} not supported!")
 if   v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!")
 elif v <= V('2.1.1'): x = 'cu{}{}-torch211'
 elif v <= V('2.1.2'): x = 'cu{}{}-torch212'
@@ -28,6 +28,7 @@
 elif v  < V('2.5.1'): x = 'cu{}{}-torch250'
 elif v <= V('2.5.1'): x = 'cu{}{}-torch251'
 elif v  < V('2.7.0'): x = 'cu{}{}-torch260'
+elif v  < V('2.8.0'): x = 'cu{}{}-torch270'
 else: raise RuntimeError(f"Torch = {v} too new!")
 x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
 print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"')