From 6f6a7a5e9bb67ecd23e6486bf79160b485c35d82 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Fri, 28 Mar 2025 16:49:12 -0700 Subject: [PATCH 01/38] add model registry --- tests/__init__.py | 0 tests/test_model_registry.py | 86 ++++++++ tests/utils/hf_hub.py | 72 +++++++ unsloth/model_registry.py | 390 +++++++++++++++++++++++++++++++++++ 4 files changed, 548 insertions(+) create mode 100644 tests/__init__.py create mode 100644 tests/test_model_registry.py create mode 100644 tests/utils/hf_hub.py create mode 100644 unsloth/model_registry.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py new file mode 100644 index 0000000000..c3eb4b0c8d --- /dev/null +++ b/tests/test_model_registry.py @@ -0,0 +1,86 @@ +from dataclasses import dataclass + +import pytest +from huggingface_hub import ModelInfo as HfModelInfo +from unsloth.model_registry import ( + ModelInfo, + get_llama_models, + get_llama_vision_models, + get_phi_instruct_models, + get_phi_models, + get_qwen_models, + get_qwen_vl_models, +) + +from .utils.hf_hub import get_model_info + +MODEL_NAMES = [ + "llama", + "llama_vision", + "qwen", + "qwen_vl", + "phi", + "phi_instruct", +] +REGISTERED_MODELS = [ + get_llama_models(), + get_llama_vision_models(), + get_qwen_models(), + get_qwen_vl_models(), + get_phi_models(), + get_phi_instruct_models(), +] + + +@dataclass +class ModelTestParam: + name: str + models: dict[str, ModelInfo] + + +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, REGISTERED_MODELS) +] + + +@pytest.mark.parametrize( + "model_test_param", TestParams, ids=lambda param: param.name +) +def test_model_uploaded(model_test_param: ModelTestParam): + missing_models = _test_model_uploaded(model_test_param.models) + assert not missing_models, ( + f"{model_test_param.name} missing following models: {missing_models}" + ) + + +if __name__ == "__main__": + for method in [ + get_llama_models, + get_llama_vision_models, + get_qwen_models, + get_qwen_vl_models, + get_phi_models, + get_phi_instruct_models, + ]: + models = method() + model_name = next(iter(models.values())).base_name + print(f"{model_name}: {len(models)} registered") + for model_info in models.values(): + print(f" {model_info.model_path}") + missing_models = test_model_uploaded(list(models.keys())) + + if missing_models: + print("--------------------------------") + print(f"Missing models: {missing_models}") + print("--------------------------------") diff --git a/tests/utils/hf_hub.py b/tests/utils/hf_hub.py new file mode 100644 index 0000000000..e3230e6ca5 --- /dev/null +++ b/tests/utils/hf_hub.py @@ -0,0 +1,72 @@ +from huggingface_hub import HfApi, ModelInfo + +api = HfApi() + +POPULARITY_PROPERTIES = [ + "downloads", + "downloadsAllTime", + "trendingScore", + "likes", +] +THOUSAND = 1000 +MILLION = 1000000 +BILLION = 1000000000 + + +def formatted_int(value: int) -> str: + if value < THOUSAND: + return str(value) + elif value < MILLION: + return f"{float(value) / 1000:,.1f}K" + elif value < BILLION: + return f"{float(value) // 1000000:,.1f}M" + + +def get_model_info( + model_id: str, properties: list[str] = ["safetensors", "lastModified"] +) -> ModelInfo: + """ + Get the model info for a specific model. + + properties: list[str] = See https://huggingface.co/docs/huggingface_hub/api-ref/hf_hub/hf_api/model_info + Default properties: ["safetensors", "lastModified"], only retrieves minimal information. + Set to None to retrieve the full model information. + """ + try: + model_info: ModelInfo = api.model_info(model_id, expand=properties) + except Exception as e: + print(f"Error getting model info for {model_id}: {e}") + model_info = None + return model_info + + +def retrieve_models( + properties: list[str] = None, + full: bool = False, + sort: str = "downloads", + author: str = "unsloth", + search: str = None, + limit: int = 10, +) -> ModelInfo: + """ + Retrieve models from the Hugging Face Hub. + + properties: list[str] = See https://huggingface.co/docs/huggingface_hub/api-ref/hf_hub/hf_api/list_models + full: bool = Whether to retrieve the full model information, if True properties will be ignored. + sort: str = The sort order. + author: str = The author of the model. + search: str = The search query for filtering models. + + """ + if full: + properties = None + + models: list[ModelInfo] = api.list_models( + author=author, + search=search, + sort=sort, + limit=limit, + expand=properties, + full=full, + ) + return models diff --git a/unsloth/model_registry.py b/unsloth/model_registry.py new file mode 100644 index 0000000000..a322ed0dc6 --- /dev/null +++ b/unsloth/model_registry.py @@ -0,0 +1,390 @@ +from dataclasses import dataclass, field +from functools import partial +from typing import Callable, Literal + +BNB_QUANTIZED_TAG = "bnb-4bit" +UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG +INSTRUCT_TAG = "Instruct" +QUANT_TYPES = [None, "bnb", "unsloth"] + +_IS_LLAMA_REGISTERED = False +_IS_LLAMA_VISION_REGISTERED = False + +_IS_QWEN_REGISTERED = False +_IS_QWEN_VL_REGISTERED = False + +_IS_GEMMA_REGISTERED = False + +_IS_PHI_REGISTERED = False +_IS_PHI_INSTRUCT_REGISTERED = False + + +def construct_model_key(org, base_name, version, size, quant_type, instruct_tag): + key = f"{org}/{base_name}-{version}-{size}B" + if instruct_tag: + key = "-".join([key, instruct_tag]) + if quant_type: + if quant_type == "bnb": + key = "-".join([key, BNB_QUANTIZED_TAG]) + elif quant_type == "unsloth": + key = "-".join([key, UNSLOTH_DYNAMIC_QUANT_TAG]) + return key + + +@dataclass +class ModelInfo: + org: str + base_name: str + version: str + size: int + name: str = None # full model name, constructed from base_name, version, and size unless provided + is_multimodal: bool = False + instruct_tag: str = None + quant_type: Literal["bnb", "unsloth"] = None + + def __post_init__(self): + self.name = self.name or self.construct_model_name( + self.base_name, + self.version, + self.size, + self.quant_type, + self.instruct_tag, + ) + + @staticmethod + def append_instruct_tag(key: str, instruct_tag: str = None): + if instruct_tag: + key = "-".join([key, instruct_tag]) + return key + + @staticmethod + def append_quant_type(key: str, quant_type: Literal["bnb", "unsloth"] = None): + if quant_type: + if quant_type == "bnb": + key = "-".join([key, BNB_QUANTIZED_TAG]) + elif quant_type == "unsloth": + key = "-".join([key, UNSLOTH_DYNAMIC_QUANT_TAG]) + return key + + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + raise NotImplementedError("Subclass must implement this method") + + @property + def model_path( + self, + ) -> str: + return f"{self.org}/{self.name}" + + +class LlamaModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +class LlamaVisionModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}-{size}B-Vision" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +class QwenModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}{version}-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +class QwenVLModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}{version}-VL-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +class PhiModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +# Llama text only models +_LLAMA_INFO = { + "org": "meta-llama", + "base_name": "Llama", + "instruct_tags": [None, "Instruct"], + "model_versions": ["3.2", "3.1"], + "model_sizes": {"3.2": [1, 3], "3.1": [8]}, + "is_multimodal": False, + "model_info_cls": LlamaModelInfo, +} + +_LLAMA_VISION_INFO = { + "org": "meta-llama", + "base_name": "Llama", + "instruct_tags": [None, "Instruct"], + "model_versions": ["3.2"], + "model_sizes": {"3.2": [11, 90]}, + "is_multimodal": True, + "model_info_cls": LlamaVisionModelInfo, +} +# Qwen text only models +# NOTE: Qwen vision models will be registered separately +_QWEN_INFO = { + "org": "Qwen", + "base_name": "Qwen", + "instruct_tags": [None, "Instruct"], + "model_versions": ["2.5"], + "model_sizes": {"2.5": [3, 7]}, + "is_multimodal": False, + "model_info_cls": QwenModelInfo, +} + +_QWEN_VL_INFO = { + "org": "Qwen", + "base_name": "Qwen", + "instruct_tags": ["Instruct"], # No base, only instruction tuned + "model_versions": ["2.5"], + "model_sizes": {"2.5": [3, 7, 32, 72]}, + "is_multimodal": True, + "instruction_tuned_only": True, + "model_info_cls": QwenVLModelInfo, +} + +_GEMMA_INFO = { + "org": "google", + "base_name": "gemma", + "instruct_tags": ["pt", "it"], # pt = base, it = instruction tuned + "model_versions": ["3"], + "model_sizes": {"3": [1, 4, 12, 27]}, + "is_multimodal": True, +} + +_PHI_INFO = { + "org": "microsoft", + "base_name": "phi", + "model_versions": ["4"], + "model_sizes": {"4": [None]}, # -1 means only 1 size + "instruct_tags": [None], + "is_multimodal": False, + "model_info_cls": PhiModelInfo, +} + +_PHI_INSTRUCT_INFO = { + "org": "microsoft", + "base_name": "Phi", + "model_versions": ["4"], + "model_sizes": {"4": [None]}, # -1 means only 1 size + "instruct_tags": ["mini-instruct"], + "is_multimodal": False, + "model_info_cls": PhiModelInfo, +} + + +MODEL_REGISTRY = {} + + +def register_model( + model_info_cls: ModelInfo, + org: str, + base_name: str, + version: str, + size: int, + quant_type: Literal["bnb", "unsloth"] = None, + is_multimodal: bool = False, + instruct_tag: str = INSTRUCT_TAG, + name: str = None, +): + name = name or model_info_cls.construct_model_name( + base_name=base_name, + version=version, + size=size, + quant_type=quant_type, + instruct_tag=instruct_tag, + ) + key = f"{org}/{name}" + + if key in MODEL_REGISTRY: + raise ValueError(f"Model {key} already registered") + + MODEL_REGISTRY[key] = model_info_cls( + org=org, + base_name=base_name, + version=version, + size=size, + is_multimodal=is_multimodal, + instruct_tag=instruct_tag, + quant_type=quant_type, + name=name, + ) + + +def _register_models(model_info: dict): + org = model_info["org"] + base_name = model_info["base_name"] + instruct_tags = model_info["instruct_tags"] + model_versions = model_info["model_versions"] + model_sizes = model_info["model_sizes"] + is_multimodal = model_info["is_multimodal"] + model_info_cls = model_info["model_info_cls"] + + for version in model_versions: + for size in model_sizes[version]: + for instruct_tag in instruct_tags: + for quant_type in QUANT_TYPES: + _org = "unsloth" if quant_type is not None else org + register_model( + model_info_cls=model_info_cls, + org=_org, + base_name=base_name, + version=version, + size=size, + instruct_tag=instruct_tag, + quant_type=quant_type, + is_multimodal=is_multimodal, + ) + + +def register_llama_models(): + global _IS_LLAMA_REGISTERED + if _IS_LLAMA_REGISTERED: + return + _register_models(_LLAMA_INFO) + _IS_LLAMA_REGISTERED = True + + +def register_llama_vision_models(): + global _IS_LLAMA_VISION_REGISTERED + if _IS_LLAMA_VISION_REGISTERED: + return + _register_models(_LLAMA_VISION_INFO) + _IS_LLAMA_VISION_REGISTERED = True + + +def register_qwen_models(): + global _IS_QWEN_REGISTERED + if _IS_QWEN_REGISTERED: + return + + _register_models(_QWEN_INFO) + _IS_QWEN_REGISTERED = True + + +def register_qwen_vl_models(): + global _IS_QWEN_VL_REGISTERED + if _IS_QWEN_VL_REGISTERED: + return + + _register_models(_QWEN_VL_INFO) + _IS_QWEN_VL_REGISTERED = True + + +def register_gemma_models(): + global _IS_GEMMA_REGISTERED + _register_models(_GEMMA_INFO) + _IS_GEMMA_REGISTERED = True + + +def register_phi_models(): + global _IS_PHI_REGISTERED + if _IS_PHI_REGISTERED: + return + _register_models(_PHI_INFO) + _IS_PHI_REGISTERED = True + + +def register_phi_instruct_models(): + global _IS_PHI_INSTRUCT_REGISTERED + if _IS_PHI_INSTRUCT_REGISTERED: + return + + _register_models(_PHI_INSTRUCT_INFO) + _IS_PHI_INSTRUCT_REGISTERED = True + + +def _base_name_filter(model_info: ModelInfo, base_name: str): + return model_info.base_name == base_name + + +def _get_models(filter_func: Callable[[ModelInfo], bool] = _base_name_filter): + return {k: v for k, v in MODEL_REGISTRY.items() if filter_func(v)} + + +def get_llama_models(): + if not _IS_LLAMA_REGISTERED: + register_llama_models() + + return _get_models(partial(_base_name_filter, base_name=_LLAMA_INFO["base_name"])) + + +def get_llama_vision_models(): + if not _IS_LLAMA_VISION_REGISTERED: + register_llama_vision_models() + + return _get_models( + lambda model_info: model_info.base_name == _LLAMA_VISION_INFO["base_name"] + and model_info.is_multimodal + ) + + +def get_qwen_models(): + if not _IS_QWEN_REGISTERED: + register_qwen_models() + + return _get_models( + lambda model_info: model_info.base_name == _QWEN_INFO["base_name"] + ) + + +def get_qwen_vl_models(): + if not _IS_QWEN_VL_REGISTERED: + register_qwen_vl_models() + return _get_models( + lambda model_info: model_info.base_name == _QWEN_VL_INFO["base_name"] + ) + + +def get_gemma_models(): + if not _IS_GEMMA_REGISTERED: + register_gemma_models() + + return _get_models( + lambda model_info: model_info.base_name == _GEMMA_INFO["base_name"] + ) + + +def get_phi_models(): + if not _IS_PHI_REGISTERED: + register_phi_models() + return _get_models( + lambda model_info: model_info.base_name == _PHI_INFO["base_name"] + ) + + +def get_phi_instruct_models(): + if not _IS_PHI_INSTRUCT_REGISTERED: + register_phi_instruct_models() + return _get_models( + lambda model_info: model_info.base_name == _PHI_INSTRUCT_INFO["base_name"] + ) + + +if __name__ == "__main__": + register_llama_models() + for k, v in MODEL_REGISTRY.items(): + print(f"{k}: {v}") + print(v.model_path) \ No newline at end of file From f21d61ae5dc22ca0683c05071c304b63e2ae60b4 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Fri, 28 Mar 2025 16:54:38 -0700 Subject: [PATCH 02/38] move hf hub utils to unsloth/utils --- pyproject.toml | 4 ++++ tests/test_model_registry.py | 3 +-- unsloth/utils/__init__.py | 0 {tests => unsloth}/utils/hf_hub.py | 0 4 files changed, 5 insertions(+), 2 deletions(-) create mode 100644 unsloth/utils/__init__.py rename {tests => unsloth}/utils/hf_hub.py (100%) diff --git a/pyproject.toml b/pyproject.toml index 7f24aabbf2..808a956c89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,10 @@ include-package-data = false exclude = ["images*", "tests*"] [project.optional-dependencies] +dev = [ + "pytest", +] + triton = [ "triton-windows ; platform_system == 'Windows'", ] diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index c3eb4b0c8d..183edc92d5 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -11,8 +11,7 @@ from unsloth.model_registry import ( get_qwen_models, get_qwen_vl_models, ) - -from .utils.hf_hub import get_model_info +from unsloth.utils.hf_hub import get_model_info MODEL_NAMES = [ "llama", diff --git a/unsloth/utils/__init__.py b/unsloth/utils/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/utils/hf_hub.py b/unsloth/utils/hf_hub.py similarity index 100% rename from tests/utils/hf_hub.py rename to unsloth/utils/hf_hub.py From 4b3df3d214ae4e7c78dc03cc4dd1b7d39b9482c0 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 10:43:00 -0700 Subject: [PATCH 03/38] refactor global model info dicts to dataclasses --- unsloth/model_registry.py | 105 +++++++++++++++++++++++++++++--------- 1 file changed, 82 insertions(+), 23 deletions(-) diff --git a/unsloth/model_registry.py b/unsloth/model_registry.py index a322ed0dc6..bb6540b5b5 100644 --- a/unsloth/model_registry.py +++ b/unsloth/model_registry.py @@ -121,6 +121,41 @@ class PhiModelInfo(ModelInfo): key = cls.append_quant_type(key, quant_type) return key +@dataclass +class ModelMetaBase: + org: str + base_name: str + +@dataclass +class ModelMeta(ModelMetaBase): + instruct_tags: list[str] + model_version: str + model_sizes: list[str] + is_multimodal: bool + model_info_cls: type[ModelInfo] + quant_types: list[Literal[None, "bnb", "unsloth", "GGUF"]] + +@dataclass +class LlamaMetaBase(ModelMetaBase): + org: str = "meta-llama" + base_name: str = "Llama" + +@dataclass +class LlamaMeta3_1(LlamaMetaBase, ModelMeta): + instruct_tags: list[str] = [None, "Instruct"] + model_version: str = "3.1" + model_sizes: list[str] = [8] + is_multimodal: bool = False + quant_types: list[Literal[None, "bnb", "unsloth"]] = [None] + model_info_cls: type[ModelInfo] = LlamaModelInfo +@dataclass +class LlamaMeta3_2(LlamaMetaBase, ModelMeta): + instruct_tags: list[str] = [None, "Instruct"] + model_version: str = "3.2" + model_sizes: list[str] = [1, 3] + is_multimodal: bool = False + quant_types: list[Literal[None, "bnb", "unsloth"]] = [None] + model_info_cls: type[ModelInfo] = LlamaModelInfo # Llama text only models _LLAMA_INFO = { @@ -233,31 +268,55 @@ def register_model( ) -def _register_models(model_info: dict): - org = model_info["org"] - base_name = model_info["base_name"] - instruct_tags = model_info["instruct_tags"] - model_versions = model_info["model_versions"] - model_sizes = model_info["model_sizes"] - is_multimodal = model_info["is_multimodal"] - model_info_cls = model_info["model_info_cls"] +# def _register_models(model_info: dict): +# org = model_info["org"] +# base_name = model_info["base_name"] +# instruct_tags = model_info["instruct_tags"] +# model_versions = model_info["model_versions"] +# model_sizes = model_info["model_sizes"] +# is_multimodal = model_info["is_multimodal"] +# model_info_cls = model_info["model_info_cls"] - for version in model_versions: - for size in model_sizes[version]: - for instruct_tag in instruct_tags: - for quant_type in QUANT_TYPES: - _org = "unsloth" if quant_type is not None else org - register_model( - model_info_cls=model_info_cls, - org=_org, - base_name=base_name, - version=version, - size=size, - instruct_tag=instruct_tag, - quant_type=quant_type, - is_multimodal=is_multimodal, - ) +# for version in model_versions: +# for size in model_sizes[version]: +# for instruct_tag in instruct_tags: +# for quant_type in QUANT_TYPES: +# _org = "unsloth" if quant_type is not None else org +# register_model( +# model_info_cls=model_info_cls, +# org=_org, +# base_name=base_name, +# version=version, +# size=size, +# instruct_tag=instruct_tag, +# quant_type=quant_type, +# is_multimodal=is_multimodal, +# ) +def _register_models(model_meta: ModelMeta): + org = model_meta.org + base_name = model_meta.base_name + instruct_tags = model_meta.instruct_tags + model_version = model_meta.model_version + model_sizes = model_meta.model_sizes + is_multimodal = model_meta.is_multimodal + quant_types = model_meta.quant_types + model_info_cls = model_meta.model_info_cls + + for size in model_sizes: + for instruct_tag in instruct_tags: + for quant_type in quant_types: + _org = "unsloth" if quant_type is not None else org + register_model( + model_info_cls=model_info_cls, + org=_org, + base_name=base_name, + version=model_version, + size=size, + instruct_tag=instruct_tag, + quant_type=quant_type, + is_multimodal=is_multimodal, + ) def register_llama_models(): global _IS_LLAMA_REGISTERED From 410c4b4c7653ad9ba3c84554f2f11b442d32e68a Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 10:58:51 -0700 Subject: [PATCH 04/38] fix dataclass init --- unsloth/model_registry.py | 151 +++++++++++++++++++++++++------------- unsloth/utils/hf_hub.py | 8 +- 2 files changed, 105 insertions(+), 54 deletions(-) diff --git a/unsloth/model_registry.py b/unsloth/model_registry.py index bb6540b5b5..dede596414 100644 --- a/unsloth/model_registry.py +++ b/unsloth/model_registry.py @@ -19,7 +19,9 @@ _IS_PHI_REGISTERED = False _IS_PHI_INSTRUCT_REGISTERED = False -def construct_model_key(org, base_name, version, size, quant_type, instruct_tag): +def construct_model_key( + org, base_name, version, size, quant_type, instruct_tag +): key = f"{org}/{base_name}-{version}-{size}B" if instruct_tag: key = "-".join([key, instruct_tag]) @@ -58,7 +60,9 @@ class ModelInfo: return key @staticmethod - def append_quant_type(key: str, quant_type: Literal["bnb", "unsloth"] = None): + def append_quant_type( + key: str, quant_type: Literal["bnb", "unsloth"] = None + ): if quant_type: if quant_type == "bnb": key = "-".join([key, BNB_QUANTIZED_TAG]) @@ -67,7 +71,9 @@ class ModelInfo: return key @classmethod - def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): raise NotImplementedError("Subclass must implement this method") @property @@ -79,7 +85,9 @@ class ModelInfo: class LlamaModelInfo(ModelInfo): @classmethod - def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): key = f"{base_name}-{version}-{size}B" key = cls.append_instruct_tag(key, instruct_tag) key = cls.append_quant_type(key, quant_type) @@ -88,7 +96,9 @@ class LlamaModelInfo(ModelInfo): class LlamaVisionModelInfo(ModelInfo): @classmethod - def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): key = f"{base_name}-{version}-{size}B-Vision" key = cls.append_instruct_tag(key, instruct_tag) key = cls.append_quant_type(key, quant_type) @@ -97,7 +107,9 @@ class LlamaVisionModelInfo(ModelInfo): class QwenModelInfo(ModelInfo): @classmethod - def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): key = f"{base_name}{version}-{size}B" key = cls.append_instruct_tag(key, instruct_tag) key = cls.append_quant_type(key, quant_type) @@ -106,7 +118,9 @@ class QwenModelInfo(ModelInfo): class QwenVLModelInfo(ModelInfo): @classmethod - def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): key = f"{base_name}{version}-VL-{size}B" key = cls.append_instruct_tag(key, instruct_tag) key = cls.append_quant_type(key, quant_type) @@ -115,58 +129,62 @@ class QwenVLModelInfo(ModelInfo): class PhiModelInfo(ModelInfo): @classmethod - def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): key = f"{base_name}-{version}" key = cls.append_instruct_tag(key, instruct_tag) key = cls.append_quant_type(key, quant_type) return key + @dataclass -class ModelMetaBase: +class ModelMeta: org: str base_name: str - -@dataclass -class ModelMeta(ModelMetaBase): - instruct_tags: list[str] model_version: str - model_sizes: list[str] - is_multimodal: bool model_info_cls: type[ModelInfo] - quant_types: list[Literal[None, "bnb", "unsloth", "GGUF"]] - -@dataclass -class LlamaMetaBase(ModelMetaBase): - org: str = "meta-llama" - base_name: str = "Llama" - -@dataclass -class LlamaMeta3_1(LlamaMetaBase, ModelMeta): - instruct_tags: list[str] = [None, "Instruct"] - model_version: str = "3.1" - model_sizes: list[str] = [8] + model_sizes: list[str] = field(default_factory=list) + instruct_tags: list[str] = field(default_factory=list) + quant_types: list[Literal[None, "bnb", "unsloth"]] = field( + default_factory=list + ) is_multimodal: bool = False - quant_types: list[Literal[None, "bnb", "unsloth"]] = [None] - model_info_cls: type[ModelInfo] = LlamaModelInfo -@dataclass -class LlamaMeta3_2(LlamaMetaBase, ModelMeta): - instruct_tags: list[str] = [None, "Instruct"] - model_version: str = "3.2" - model_sizes: list[str] = [1, 3] - is_multimodal: bool = False - quant_types: list[Literal[None, "bnb", "unsloth"]] = [None] - model_info_cls: type[ModelInfo] = LlamaModelInfo -# Llama text only models -_LLAMA_INFO = { - "org": "meta-llama", - "base_name": "Llama", - "instruct_tags": [None, "Instruct"], - "model_versions": ["3.2", "3.1"], - "model_sizes": {"3.2": [1, 3], "3.1": [8]}, - "is_multimodal": False, - "model_info_cls": LlamaModelInfo, -} + +LlamaMeta3_1 = ModelMeta( + org="meta-llama", + base_name="Llama", + instruct_tags=[None, "Instruct"], + model_version="3.1", + model_sizes=[8], + model_info_cls=LlamaModelInfo, + is_multimodal=False, + quant_types=[None, "bnb", "unsloth"], +) + +LlamaMeta3_2 = ModelMeta( + org="meta-llama", + base_name="Llama", + instruct_tags=[None, "Instruct"], + model_version="3.2", + model_sizes=[1, 3], + model_info_cls=LlamaModelInfo, + is_multimodal=False, + quant_types=[None, "bnb", "unsloth"], +) + + +# # Llama text only models +# _LLAMA_INFO = { +# "org": "meta-llama", +# "base_name": "Llama", +# "instruct_tags": [None, "Instruct"], +# "model_versions": ["3.2", "3.1"], +# "model_sizes": {"3.2": [1, 3], "3.1": [8]}, +# "is_multimodal": False, +# "model_info_cls": LlamaModelInfo, +# } _LLAMA_VISION_INFO = { "org": "meta-llama", @@ -293,6 +311,7 @@ def register_model( # is_multimodal=is_multimodal, # ) + def _register_models(model_meta: ModelMeta): org = model_meta.org base_name = model_meta.base_name @@ -318,6 +337,7 @@ def _register_models(model_meta: ModelMeta): is_multimodal=is_multimodal, ) + def register_llama_models(): global _IS_LLAMA_REGISTERED if _IS_LLAMA_REGISTERED: @@ -387,7 +407,9 @@ def get_llama_models(): if not _IS_LLAMA_REGISTERED: register_llama_models() - return _get_models(partial(_base_name_filter, base_name=_LLAMA_INFO["base_name"])) + return _get_models( + partial(_base_name_filter, base_name=_LLAMA_INFO["base_name"]) + ) def get_llama_vision_models(): @@ -395,7 +417,8 @@ def get_llama_vision_models(): register_llama_vision_models() return _get_models( - lambda model_info: model_info.base_name == _LLAMA_VISION_INFO["base_name"] + lambda model_info: model_info.base_name + == _LLAMA_VISION_INFO["base_name"] and model_info.is_multimodal ) @@ -438,12 +461,34 @@ def get_phi_instruct_models(): if not _IS_PHI_INSTRUCT_REGISTERED: register_phi_instruct_models() return _get_models( - lambda model_info: model_info.base_name == _PHI_INSTRUCT_INFO["base_name"] + lambda model_info: model_info.base_name + == _PHI_INSTRUCT_INFO["base_name"] ) if __name__ == "__main__": - register_llama_models() + from huggingface_hub import HfApi + + api = HfApi() + + def get_model_info( + model_id: str, properties: list[str] = None + ) -> ModelInfo: + try: + model_info: ModelInfo = api.model_info(model_id, expand=properties) + except Exception as e: + print(f"Error getting model info for {model_id}: {e}") + model_info = None + return model_info + + test_model = LlamaMeta3_2 + _register_models(test_model) + for k, v in MODEL_REGISTRY.items(): - print(f"{k}: {v}") - print(v.model_path) \ No newline at end of file + model_info = get_model_info(v.model_path) + if model_info is None: + # print unicode cross mark followed by model k + print(f"\u2718 {k}") + else: + # print unicode checkmark followed by model k + print(f"\u2713 {k} found") diff --git a/unsloth/utils/hf_hub.py b/unsloth/utils/hf_hub.py index e3230e6ca5..da3f72a18e 100644 --- a/unsloth/utils/hf_hub.py +++ b/unsloth/utils/hf_hub.py @@ -1,6 +1,6 @@ from huggingface_hub import HfApi, ModelInfo -api = HfApi() +api: HfApi POPULARITY_PROPERTIES = [ "downloads", @@ -32,6 +32,9 @@ def get_model_info( Default properties: ["safetensors", "lastModified"], only retrieves minimal information. Set to None to retrieve the full model information. """ + global api + if api is None: + api = HfApi() try: model_info: ModelInfo = api.model_info(model_id, expand=properties) except Exception as e: @@ -58,6 +61,9 @@ def retrieve_models( search: str = The search query for filtering models. """ + global api + if api is None: + api = HfApi() if full: properties = None From 85209602f31f0ea6f31bfdb6ac88fefe1d37a4bb Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 11:06:11 -0700 Subject: [PATCH 05/38] fix llama registration --- unsloth/model_registry.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/unsloth/model_registry.py b/unsloth/model_registry.py index dede596414..2f7ccb956d 100644 --- a/unsloth/model_registry.py +++ b/unsloth/model_registry.py @@ -342,7 +342,8 @@ def register_llama_models(): global _IS_LLAMA_REGISTERED if _IS_LLAMA_REGISTERED: return - _register_models(_LLAMA_INFO) + _register_models(LlamaMeta3_1) + _register_models(LlamaMeta3_2) _IS_LLAMA_REGISTERED = True @@ -403,13 +404,18 @@ def _get_models(filter_func: Callable[[ModelInfo], bool] = _base_name_filter): return {k: v for k, v in MODEL_REGISTRY.items() if filter_func(v)} -def get_llama_models(): +def get_llama_models(version: str = None): if not _IS_LLAMA_REGISTERED: register_llama_models() - return _get_models( - partial(_base_name_filter, base_name=_LLAMA_INFO["base_name"]) + llama_models: dict[str, ModelInfo] = _get_models( + partial(_base_name_filter, base_name=LlamaMeta3_1.base_name) ) + if version is not None: + llama_models = { + k: v for k, v in llama_models.items() if v.version == version + } + return llama_models def get_llama_vision_models(): @@ -481,14 +487,17 @@ if __name__ == "__main__": model_info = None return model_info - test_model = LlamaMeta3_2 - _register_models(test_model) + register_llama_models() - for k, v in MODEL_REGISTRY.items(): + llama3_1_models = get_llama_models(version="3.2") + missing_models = [] + for k, v in llama3_1_models.items(): model_info = get_model_info(v.model_path) if model_info is None: # print unicode cross mark followed by model k print(f"\u2718 {k}") - else: - # print unicode checkmark followed by model k - print(f"\u2713 {k} found") + missing_models.append(k) + + if len(missing_models) == 0: + # print unicode checkmark + print(f"\u2713 All models found!") \ No newline at end of file From 35e3b48c2b751a7c55728fed11c3a8829f0bd452 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 11:06:59 -0700 Subject: [PATCH 06/38] remove deprecated key function --- unsloth/model_registry.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/unsloth/model_registry.py b/unsloth/model_registry.py index 2f7ccb956d..dfdf3755ed 100644 --- a/unsloth/model_registry.py +++ b/unsloth/model_registry.py @@ -19,20 +19,6 @@ _IS_PHI_REGISTERED = False _IS_PHI_INSTRUCT_REGISTERED = False -def construct_model_key( - org, base_name, version, size, quant_type, instruct_tag -): - key = f"{org}/{base_name}-{version}-{size}B" - if instruct_tag: - key = "-".join([key, instruct_tag]) - if quant_type: - if quant_type == "bnb": - key = "-".join([key, BNB_QUANTIZED_TAG]) - elif quant_type == "unsloth": - key = "-".join([key, UNSLOTH_DYNAMIC_QUANT_TAG]) - return key - - @dataclass class ModelInfo: org: str From ab7c51b4a5340cc6c25fe0b31784b44bd3de26e8 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 11:36:48 -0700 Subject: [PATCH 07/38] start registry reog --- .gitignore | 177 ++++++++++++++ unsloth/registry/__init__.py | 0 unsloth/registry/_llama.py | 77 +++++++ unsloth/{ => registry}/model_registry.py | 279 ++++++----------------- unsloth/registry/registry.py | 149 ++++++++++++ 5 files changed, 478 insertions(+), 204 deletions(-) create mode 100644 .gitignore create mode 100644 unsloth/registry/__init__.py create mode 100644 unsloth/registry/_llama.py rename unsloth/{ => registry}/model_registry.py (54%) create mode 100644 unsloth/registry/registry.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000..ceb66ed122 --- /dev/null +++ b/.gitignore @@ -0,0 +1,177 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# unsloth compiled cache +unsloth_compiled_cache diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py new file mode 100644 index 0000000000..35b40dccb9 --- /dev/null +++ b/unsloth/registry/_llama.py @@ -0,0 +1,77 @@ +from unsloth.registry.registry import ModelInfo, ModelMeta, _register_models + +_IS_LLAMA_REGISTERED = False + +class LlamaModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +class LlamaVisionModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}-{size}B-Vision" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +# Llama 3.1 +LlamaMeta3_1 = ModelMeta( + org="meta-llama", + base_name="Llama", + instruct_tags=[None, "Instruct"], + model_version="3.1", + model_sizes=[8], + model_info_cls=LlamaModelInfo, + is_multimodal=False, + quant_types=[None, "bnb", "unsloth"], +) + +# Llama 3.2 +LlamaMeta3_2 = ModelMeta( + org="meta-llama", + base_name="Llama", + instruct_tags=[None, "Instruct"], + model_version="3.2", + model_sizes=[1, 3], + model_info_cls=LlamaModelInfo, + is_multimodal=False, + quant_types=[None, "bnb", "unsloth"], +) + +# Llama 3.2 Vision +LlamaMeta3_2_Vision = ModelMeta( + org="meta-llama", + base_name="Llama", + instruct_tags=[None, "Instruct"], + model_version="3.2", + model_sizes=[11, 90], + model_info_cls=LlamaVisionModelInfo, + is_multimodal=True, + quant_types=[None, "bnb", "unsloth"], +) + + +def register_llama_models(): + global _IS_LLAMA_REGISTERED + if _IS_LLAMA_REGISTERED: + return + _register_models(LlamaMeta3_1) + _register_models(LlamaMeta3_2) + _IS_LLAMA_REGISTERED = True + +register_llama_models() + +if __name__ == "__main__": + from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + for model_id, model_info in MODEL_REGISTRY.items(): + model_info = _check_model_info(model_id) + if model_info is None: + print(f"\u2718 {model_id}") + else: + print(f"\u2713 {model_id}") \ No newline at end of file diff --git a/unsloth/model_registry.py b/unsloth/registry/model_registry.py similarity index 54% rename from unsloth/model_registry.py rename to unsloth/registry/model_registry.py index dfdf3755ed..a0cd71c17a 100644 --- a/unsloth/model_registry.py +++ b/unsloth/registry/model_registry.py @@ -1,11 +1,8 @@ -from dataclasses import dataclass, field from functools import partial from typing import Callable, Literal -BNB_QUANTIZED_TAG = "bnb-4bit" -UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG -INSTRUCT_TAG = "Instruct" -QUANT_TYPES = [None, "bnb", "unsloth"] +from unsloth.registry._llama import LlamaMeta3_1, LlamaMeta3_2 +from unsloth.registry.common import ModelInfo, ModelMeta _IS_LLAMA_REGISTERED = False _IS_LLAMA_VISION_REGISTERED = False @@ -19,222 +16,97 @@ _IS_PHI_REGISTERED = False _IS_PHI_INSTRUCT_REGISTERED = False -@dataclass -class ModelInfo: - org: str - base_name: str - version: str - size: int - name: str = None # full model name, constructed from base_name, version, and size unless provided - is_multimodal: bool = False - instruct_tag: str = None - quant_type: Literal["bnb", "unsloth"] = None - def __post_init__(self): - self.name = self.name or self.construct_model_name( - self.base_name, - self.version, - self.size, - self.quant_type, - self.instruct_tag, - ) - - @staticmethod - def append_instruct_tag(key: str, instruct_tag: str = None): - if instruct_tag: - key = "-".join([key, instruct_tag]) - return key - - @staticmethod - def append_quant_type( - key: str, quant_type: Literal["bnb", "unsloth"] = None - ): - if quant_type: - if quant_type == "bnb": - key = "-".join([key, BNB_QUANTIZED_TAG]) - elif quant_type == "unsloth": - key = "-".join([key, UNSLOTH_DYNAMIC_QUANT_TAG]) - return key - - @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): - raise NotImplementedError("Subclass must implement this method") - - @property - def model_path( - self, - ) -> str: - return f"{self.org}/{self.name}" +# class QwenModelInfo(ModelInfo): +# @classmethod +# def construct_model_name( +# cls, base_name, version, size, quant_type, instruct_tag +# ): +# key = f"{base_name}{version}-{size}B" +# key = cls.append_instruct_tag(key, instruct_tag) +# key = cls.append_quant_type(key, quant_type) +# return key -class LlamaModelInfo(ModelInfo): - @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): - key = f"{base_name}-{version}-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key +# class QwenVLModelInfo(ModelInfo): +# @classmethod +# def construct_model_name( +# cls, base_name, version, size, quant_type, instruct_tag +# ): +# key = f"{base_name}{version}-VL-{size}B" +# key = cls.append_instruct_tag(key, instruct_tag) +# key = cls.append_quant_type(key, quant_type) +# return key -class LlamaVisionModelInfo(ModelInfo): - @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): - key = f"{base_name}-{version}-{size}B-Vision" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key +# class PhiModelInfo(ModelInfo): +# @classmethod +# def construct_model_name( +# cls, base_name, version, size, quant_type, instruct_tag +# ): +# key = f"{base_name}-{version}" +# key = cls.append_instruct_tag(key, instruct_tag) +# key = cls.append_quant_type(key, quant_type) +# return key -class QwenModelInfo(ModelInfo): - @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): - key = f"{base_name}{version}-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key -class QwenVLModelInfo(ModelInfo): - @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): - key = f"{base_name}{version}-VL-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key - -class PhiModelInfo(ModelInfo): - @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): - key = f"{base_name}-{version}" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key - - -@dataclass -class ModelMeta: - org: str - base_name: str - model_version: str - model_info_cls: type[ModelInfo] - model_sizes: list[str] = field(default_factory=list) - instruct_tags: list[str] = field(default_factory=list) - quant_types: list[Literal[None, "bnb", "unsloth"]] = field( - default_factory=list - ) - is_multimodal: bool = False - - -LlamaMeta3_1 = ModelMeta( - org="meta-llama", - base_name="Llama", - instruct_tags=[None, "Instruct"], - model_version="3.1", - model_sizes=[8], - model_info_cls=LlamaModelInfo, - is_multimodal=False, - quant_types=[None, "bnb", "unsloth"], -) - -LlamaMeta3_2 = ModelMeta( - org="meta-llama", - base_name="Llama", - instruct_tags=[None, "Instruct"], - model_version="3.2", - model_sizes=[1, 3], - model_info_cls=LlamaModelInfo, - is_multimodal=False, - quant_types=[None, "bnb", "unsloth"], -) - - -# # Llama text only models -# _LLAMA_INFO = { -# "org": "meta-llama", -# "base_name": "Llama", +# # Qwen text only models +# # NOTE: Qwen vision models will be registered separately +# _QWEN_INFO = { +# "org": "Qwen", +# "base_name": "Qwen", # "instruct_tags": [None, "Instruct"], -# "model_versions": ["3.2", "3.1"], -# "model_sizes": {"3.2": [1, 3], "3.1": [8]}, +# "model_versions": ["2.5"], +# "model_sizes": {"2.5": [3, 7]}, # "is_multimodal": False, -# "model_info_cls": LlamaModelInfo, +# "model_info_cls": QwenModelInfo, # } -_LLAMA_VISION_INFO = { - "org": "meta-llama", - "base_name": "Llama", - "instruct_tags": [None, "Instruct"], - "model_versions": ["3.2"], - "model_sizes": {"3.2": [11, 90]}, - "is_multimodal": True, - "model_info_cls": LlamaVisionModelInfo, -} -# Qwen text only models -# NOTE: Qwen vision models will be registered separately -_QWEN_INFO = { - "org": "Qwen", - "base_name": "Qwen", - "instruct_tags": [None, "Instruct"], - "model_versions": ["2.5"], - "model_sizes": {"2.5": [3, 7]}, - "is_multimodal": False, - "model_info_cls": QwenModelInfo, -} +# _QWEN_VL_INFO = { +# "org": "Qwen", +# "base_name": "Qwen", +# "instruct_tags": ["Instruct"], # No base, only instruction tuned +# "model_versions": ["2.5"], +# "model_sizes": {"2.5": [3, 7, 32, 72]}, +# "is_multimodal": True, +# "instruction_tuned_only": True, +# "model_info_cls": QwenVLModelInfo, +# } -_QWEN_VL_INFO = { - "org": "Qwen", - "base_name": "Qwen", - "instruct_tags": ["Instruct"], # No base, only instruction tuned - "model_versions": ["2.5"], - "model_sizes": {"2.5": [3, 7, 32, 72]}, - "is_multimodal": True, - "instruction_tuned_only": True, - "model_info_cls": QwenVLModelInfo, -} +# _GEMMA_INFO = { +# "org": "google", +# "base_name": "gemma", +# "instruct_tags": ["pt", "it"], # pt = base, it = instruction tuned +# "model_versions": ["3"], +# "model_sizes": {"3": [1, 4, 12, 27]}, +# "is_multimodal": True, +# } -_GEMMA_INFO = { - "org": "google", - "base_name": "gemma", - "instruct_tags": ["pt", "it"], # pt = base, it = instruction tuned - "model_versions": ["3"], - "model_sizes": {"3": [1, 4, 12, 27]}, - "is_multimodal": True, -} +# _PHI_INFO = { +# "org": "microsoft", +# "base_name": "phi", +# "model_versions": ["4"], +# "model_sizes": {"4": [None]}, # -1 means only 1 size +# "instruct_tags": [None], +# "is_multimodal": False, +# "model_info_cls": PhiModelInfo, +# } -_PHI_INFO = { - "org": "microsoft", - "base_name": "phi", - "model_versions": ["4"], - "model_sizes": {"4": [None]}, # -1 means only 1 size - "instruct_tags": [None], - "is_multimodal": False, - "model_info_cls": PhiModelInfo, -} - -_PHI_INSTRUCT_INFO = { - "org": "microsoft", - "base_name": "Phi", - "model_versions": ["4"], - "model_sizes": {"4": [None]}, # -1 means only 1 size - "instruct_tags": ["mini-instruct"], - "is_multimodal": False, - "model_info_cls": PhiModelInfo, -} +# _PHI_INSTRUCT_INFO = { +# "org": "microsoft", +# "base_name": "Phi", +# "model_versions": ["4"], +# "model_sizes": {"4": [None]}, # -1 means only 1 size +# "instruct_tags": ["mini-instruct"], +# "is_multimodal": False, +# "model_info_cls": PhiModelInfo, +# } -MODEL_REGISTRY = {} +MODEL_REGISTRY: dict[str, ModelInfo] = {} def register_model( @@ -243,9 +115,9 @@ def register_model( base_name: str, version: str, size: int, + instruct_tag: str = None, quant_type: Literal["bnb", "unsloth"] = None, is_multimodal: bool = False, - instruct_tag: str = INSTRUCT_TAG, name: str = None, ): name = name or model_info_cls.construct_model_name( @@ -323,7 +195,6 @@ def _register_models(model_meta: ModelMeta): is_multimodal=is_multimodal, ) - def register_llama_models(): global _IS_LLAMA_REGISTERED if _IS_LLAMA_REGISTERED: diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py new file mode 100644 index 0000000000..172b6e8e86 --- /dev/null +++ b/unsloth/registry/registry.py @@ -0,0 +1,149 @@ +from dataclasses import dataclass, field +from typing import Literal + +BNB_QUANTIZED_TAG = "bnb-4bit" +UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG +QUANT_TYPE_MAP = { + "bnb": BNB_QUANTIZED_TAG, + "unsloth": UNSLOTH_DYNAMIC_QUANT_TAG, + "GGUF": "GGUF", +} +QUANT_TYPES = list(QUANT_TYPE_MAP.keys()) + + +@dataclass +class ModelInfo: + org: str + base_name: str + version: str + size: int + name: str = None # full model name, constructed from base_name, version, and size unless provided + is_multimodal: bool = False + instruct_tag: str = None + quant_type: Literal["bnb", "unsloth"] = None + + def __post_init__(self): + self.name = self.name or self.construct_model_name( + self.base_name, + self.version, + self.size, + self.quant_type, + self.instruct_tag, + ) + + @staticmethod + def append_instruct_tag(key: str, instruct_tag: str = None): + if instruct_tag: + key = "-".join([key, instruct_tag]) + return key + + @staticmethod + def append_quant_type( + key: str, quant_type: Literal["bnb", "unsloth", "GGUF"] = None + ): + if quant_type: + if quant_type == "bnb": + key = "-".join([key, QUANT_TYPE_MAP["bnb"]]) + elif quant_type == "unsloth": + key = "-".join([key, QUANT_TYPE_MAP["unsloth"]]) + elif quant_type == "GGUF": + key = "-".join([key, QUANT_TYPE_MAP["GGUF"]]) + return key + + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + raise NotImplementedError("Subclass must implement this method") + + @property + def model_path( + self, + ) -> str: + return f"{self.org}/{self.name}" + + +@dataclass +class ModelMeta: + org: str + base_name: str + model_version: str + model_info_cls: type[ModelInfo] + model_sizes: list[str] = field(default_factory=list) + instruct_tags: list[str] = field(default_factory=list) + quant_types: list[Literal[None, "bnb", "unsloth"]] = field(default_factory=list) + is_multimodal: bool = False + + +MODEL_REGISTRY: dict[str, ModelInfo] = {} + + +def register_model( + model_info_cls: ModelInfo, + org: str, + base_name: str, + version: str, + size: int, + instruct_tag: str = None, + quant_type: Literal["bnb", "unsloth"] = None, + is_multimodal: bool = False, + name: str = None, +): + name = name or model_info_cls.construct_model_name( + base_name=base_name, + version=version, + size=size, + quant_type=quant_type, + instruct_tag=instruct_tag, + ) + key = f"{org}/{name}" + + if key in MODEL_REGISTRY: + raise ValueError(f"Model {key} already registered") + + MODEL_REGISTRY[key] = model_info_cls( + org=org, + base_name=base_name, + version=version, + size=size, + is_multimodal=is_multimodal, + instruct_tag=instruct_tag, + quant_type=quant_type, + name=name, + ) + +def _check_model_info(model_id: str, properties: list[str] = ["lastModified"]): + from huggingface_hub import HfApi + from huggingface_hub import ModelInfo as HfModelInfo + api = HfApi() + + try: + model_info: HfModelInfo = api.model_info(model_id, expand=properties) + except Exception as e: + print(f"Error getting model info for {model_id}: {e}") + model_info = None + return model_info + + +def _register_models(model_meta: ModelMeta): + org = model_meta.org + base_name = model_meta.base_name + instruct_tags = model_meta.instruct_tags + model_version = model_meta.model_version + model_sizes = model_meta.model_sizes + is_multimodal = model_meta.is_multimodal + quant_types = model_meta.quant_types + model_info_cls = model_meta.model_info_cls + + for size in model_sizes: + for instruct_tag in instruct_tags: + for quant_type in quant_types: + _org = "unsloth" if quant_type is not None else org + register_model( + model_info_cls=model_info_cls, + org=_org, + base_name=base_name, + version=model_version, + size=size, + instruct_tag=instruct_tag, + quant_type=quant_type, + is_multimodal=is_multimodal, + ) From 3d1249a551afb815a60603d528d32b148642ad17 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 11:44:52 -0700 Subject: [PATCH 08/38] add llama vision --- unsloth/registry/_llama.py | 10 ++++++++++ unsloth/registry/registry.py | 9 +++++++-- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index 35b40dccb9..f1d5f6da3e 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -1,6 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, _register_models _IS_LLAMA_REGISTERED = False +_IS_LLAMA_VISION_REGISTERED = False class LlamaModelInfo(ModelInfo): @classmethod @@ -65,7 +66,16 @@ def register_llama_models(): _register_models(LlamaMeta3_2) _IS_LLAMA_REGISTERED = True + +def register_llama_vision_models(): + global _IS_LLAMA_VISION_REGISTERED + if _IS_LLAMA_VISION_REGISTERED: + return + _register_models(LlamaMeta3_2_Vision) + _IS_LLAMA_VISION_REGISTERED = True + register_llama_models() +register_llama_vision_models() if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 172b6e8e86..2402282d6a 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -113,13 +113,18 @@ def register_model( def _check_model_info(model_id: str, properties: list[str] = ["lastModified"]): from huggingface_hub import HfApi from huggingface_hub import ModelInfo as HfModelInfo + from huggingface_hub.utils import RepositoryNotFoundError api = HfApi() try: model_info: HfModelInfo = api.model_info(model_id, expand=properties) except Exception as e: - print(f"Error getting model info for {model_id}: {e}") - model_info = None + + if isinstance(e, RepositoryNotFoundError): + print(f"\u2718 {model_id} not found") + model_info = None + else: + raise e return model_info From 6b4bf12873ea7648b7d67ba16ff6959c8a88917c Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 14:37:30 -0700 Subject: [PATCH 09/38] quant types -> Enum --- unsloth/registry/registry.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 2402282d6a..2ea0618127 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -1,12 +1,23 @@ from dataclasses import dataclass, field +from enum import Enum from typing import Literal + +class QuantType(Enum): + BNB = "bnb" + UNSLOTH = "unsloth" + GGUF = "GGUF" + NONE = "none" + BNB_QUANTIZED_TAG = "bnb-4bit" UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG +GGUF_TAG = "GGUF" + QUANT_TYPE_MAP = { - "bnb": BNB_QUANTIZED_TAG, - "unsloth": UNSLOTH_DYNAMIC_QUANT_TAG, - "GGUF": "GGUF", + QuantType.BNB: BNB_QUANTIZED_TAG, + QuantType.UNSLOTH: UNSLOTH_DYNAMIC_QUANT_TAG, + QuantType.GGUF: GGUF_TAG, + QuantType.NONE: None, } QUANT_TYPES = list(QUANT_TYPE_MAP.keys()) @@ -110,16 +121,17 @@ def register_model( name=name, ) + def _check_model_info(model_id: str, properties: list[str] = ["lastModified"]): from huggingface_hub import HfApi from huggingface_hub import ModelInfo as HfModelInfo from huggingface_hub.utils import RepositoryNotFoundError + api = HfApi() try: model_info: HfModelInfo = api.model_info(model_id, expand=properties) except Exception as e: - if isinstance(e, RepositoryNotFoundError): print(f"\u2718 {model_id} not found") model_info = None From 6abdb1fef68b1c0dfd92becc857c408b813c9a4e Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 14:39:57 -0700 Subject: [PATCH 10/38] remap literal quant types to QuantType Enum --- unsloth/registry/registry.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 2ea0618127..bac0a2697e 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -31,7 +31,7 @@ class ModelInfo: name: str = None # full model name, constructed from base_name, version, and size unless provided is_multimodal: bool = False instruct_tag: str = None - quant_type: Literal["bnb", "unsloth"] = None + quant_type: QuantType = None def __post_init__(self): self.name = self.name or self.construct_model_name( @@ -50,7 +50,7 @@ class ModelInfo: @staticmethod def append_quant_type( - key: str, quant_type: Literal["bnb", "unsloth", "GGUF"] = None + key: str, quant_type: QuantType = None ): if quant_type: if quant_type == "bnb": @@ -80,7 +80,7 @@ class ModelMeta: model_info_cls: type[ModelInfo] model_sizes: list[str] = field(default_factory=list) instruct_tags: list[str] = field(default_factory=list) - quant_types: list[Literal[None, "bnb", "unsloth"]] = field(default_factory=list) + quant_types: list[QuantType] = field(default_factory=list) is_multimodal: bool = False @@ -94,7 +94,7 @@ def register_model( version: str, size: int, instruct_tag: str = None, - quant_type: Literal["bnb", "unsloth"] = None, + quant_type: QuantType = None, is_multimodal: bool = False, name: str = None, ): From 0130265ca8719aa612f4a222f2686db5b29fe891 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 15:05:33 -0700 Subject: [PATCH 11/38] add llama model registration --- unsloth/registry/_llama.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index f1d5f6da3e..211b3ac89e 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -1,4 +1,4 @@ -from unsloth.registry.registry import ModelInfo, ModelMeta, _register_models +from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models _IS_LLAMA_REGISTERED = False _IS_LLAMA_VISION_REGISTERED = False @@ -30,7 +30,7 @@ LlamaMeta3_1 = ModelMeta( model_sizes=[8], model_info_cls=LlamaModelInfo, is_multimodal=False, - quant_types=[None, "bnb", "unsloth"], + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], ) # Llama 3.2 @@ -42,7 +42,7 @@ LlamaMeta3_2 = ModelMeta( model_sizes=[1, 3], model_info_cls=LlamaModelInfo, is_multimodal=False, - quant_types=[None, "bnb", "unsloth"], + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], ) # Llama 3.2 Vision @@ -54,7 +54,7 @@ LlamaMeta3_2_Vision = ModelMeta( model_sizes=[11, 90], model_info_cls=LlamaVisionModelInfo, is_multimodal=True, - quant_types=[None, "bnb", "unsloth"], + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], ) From 95da062046a944d9df61bae48deed49f1bee7c0c Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 16:14:33 -0700 Subject: [PATCH 12/38] fix quant tag mapping --- unsloth/registry/registry.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index bac0a2697e..d045f5bd55 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -13,13 +13,12 @@ BNB_QUANTIZED_TAG = "bnb-4bit" UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG GGUF_TAG = "GGUF" -QUANT_TYPE_MAP = { +QUANT_TAG_MAP = { QuantType.BNB: BNB_QUANTIZED_TAG, QuantType.UNSLOTH: UNSLOTH_DYNAMIC_QUANT_TAG, QuantType.GGUF: GGUF_TAG, QuantType.NONE: None, } -QUANT_TYPES = list(QUANT_TYPE_MAP.keys()) @dataclass @@ -52,13 +51,8 @@ class ModelInfo: def append_quant_type( key: str, quant_type: QuantType = None ): - if quant_type: - if quant_type == "bnb": - key = "-".join([key, QUANT_TYPE_MAP["bnb"]]) - elif quant_type == "unsloth": - key = "-".join([key, QUANT_TYPE_MAP["unsloth"]]) - elif quant_type == "GGUF": - key = "-".join([key, QUANT_TYPE_MAP["GGUF"]]) + if quant_type != QuantType.NONE: + key = "-".join([key, QUANT_TAG_MAP[quant_type]]) return key @classmethod @@ -108,7 +102,7 @@ def register_model( key = f"{org}/{name}" if key in MODEL_REGISTRY: - raise ValueError(f"Model {key} already registered") + raise ValueError(f"Model {key} already registered, current keys: {MODEL_REGISTRY.keys()}") MODEL_REGISTRY[key] = model_info_cls( org=org, From 0395604928bcff4b06a9559476d90c05b76586ce Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 08:45:53 -0700 Subject: [PATCH 13/38] add qwen2.5 models to registry --- unsloth/registry/_qwen.py | 77 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 unsloth/registry/_qwen.py diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py new file mode 100644 index 0000000000..92f366bb76 --- /dev/null +++ b/unsloth/registry/_qwen.py @@ -0,0 +1,77 @@ +from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models + +_IS_QWEN_REGISTERED = False +_IS_QWEN_VL_REGISTERED = False + +class QwenModelInfo(ModelInfo): + @classmethod + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): + key = f"{base_name}{version}-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +class QwenVLModelInfo(ModelInfo): + @classmethod + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): + key = f"{base_name}{version}-VL-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + + +# Qwen Model Meta +QwenMeta = ModelMeta( + org="Qwen", + base_name="Qwen", + instruct_tags=[None, "Instruct"], + model_version="2.5", + model_sizes=[3, 7], + model_info_cls=QwenModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], +) + +# Qwen VL Model Meta +QwenVLMeta = ModelMeta( + org="Qwen", + base_name="Qwen", + instruct_tags=["Instruct"], # No base, only instruction tuned + model_version="2.5", + model_sizes=[3, 7, 32, 72], + model_info_cls=QwenVLModelInfo, + is_multimodal=True, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], +) + +def register_qwen_models(): + global _IS_QWEN_REGISTERED + if _IS_QWEN_REGISTERED: + return + _register_models(QwenMeta) + _IS_QWEN_REGISTERED = True + +def register_qwen_vl_models(): + global _IS_QWEN_VL_REGISTERED + if _IS_QWEN_VL_REGISTERED: + return + _register_models(QwenVLMeta) + _IS_QWEN_VL_REGISTERED = True + +register_qwen_models() +register_qwen_vl_models() + + +if __name__ == "__main__": + from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + for model_id, model_info in MODEL_REGISTRY.items(): + model_info = _check_model_info(model_id) + if model_info is None: + print(f"\u2718 {model_id}") + else: + print(f"\u2713 {model_id}") From 671dd3dc14b5f0918bffc1a1d99e0758a8f66dab Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 09:09:34 -0700 Subject: [PATCH 14/38] add option to include original model in registry --- unsloth/registry/_qwen.py | 44 +++++++++++++++++++++++++++++------- unsloth/registry/registry.py | 16 +++++++++++-- 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py index 92f366bb76..2ea340b813 100644 --- a/unsloth/registry/_qwen.py +++ b/unsloth/registry/_qwen.py @@ -2,7 +2,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register _IS_QWEN_REGISTERED = False _IS_QWEN_VL_REGISTERED = False - +_IS_QWEN_QWQ_REGISTERED = False class QwenModelInfo(ModelInfo): @classmethod def construct_model_name( @@ -24,7 +24,16 @@ class QwenVLModelInfo(ModelInfo): key = cls.append_quant_type(key, quant_type) return key - +class QwenQwQModelInfo(ModelInfo): + @classmethod + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): + key = f"{base_name}-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + # Qwen Model Meta QwenMeta = ModelMeta( org="Qwen", @@ -49,23 +58,42 @@ QwenVLMeta = ModelMeta( quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], ) -def register_qwen_models(): +# Qwen QwQ Model Meta +QwenQwQMeta = ModelMeta( + org="Qwen", + base_name="QwQ", + instruct_tags=[None], + model_version="", + model_sizes=[32], + model_info_cls=QwenQwQModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], +) + +def register_qwen_models(include_original_model: bool = False): global _IS_QWEN_REGISTERED if _IS_QWEN_REGISTERED: return - _register_models(QwenMeta) + _register_models(QwenMeta, include_original_model) _IS_QWEN_REGISTERED = True -def register_qwen_vl_models(): +def register_qwen_vl_models(include_original_model: bool = False): global _IS_QWEN_VL_REGISTERED if _IS_QWEN_VL_REGISTERED: return - _register_models(QwenVLMeta) + _register_models(QwenVLMeta, include_original_model) _IS_QWEN_VL_REGISTERED = True -register_qwen_models() -register_qwen_vl_models() +def register_qwen_qwq_models(include_original_model: bool = False): + global _IS_QWEN_QWQ_REGISTERED + if _IS_QWEN_QWQ_REGISTERED: + return + _register_models(QwenQwQMeta, include_original_model) + _IS_QWEN_QWQ_REGISTERED = True +# register_qwen_models() +# register_qwen_vl_models() +register_qwen_qwq_models(include_original_model=True) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index d045f5bd55..3ca7c20f8f 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -134,7 +134,7 @@ def _check_model_info(model_id: str, properties: list[str] = ["lastModified"]): return model_info -def _register_models(model_meta: ModelMeta): +def _register_models(model_meta: ModelMeta, include_original_model: bool = False): org = model_meta.org base_name = model_meta.base_name instruct_tags = model_meta.instruct_tags @@ -147,7 +147,7 @@ def _register_models(model_meta: ModelMeta): for size in model_sizes: for instruct_tag in instruct_tags: for quant_type in quant_types: - _org = "unsloth" if quant_type is not None else org + _org = "unsloth" # unsloth models -- these are all quantized versions of the original model register_model( model_info_cls=model_info_cls, org=_org, @@ -158,3 +158,15 @@ def _register_models(model_meta: ModelMeta): quant_type=quant_type, is_multimodal=is_multimodal, ) + # include original model from releasing organization + if include_original_model: + register_model( + model_info_cls=model_info_cls, + org=org, + base_name=base_name, + version=model_version, + size=size, + instruct_tag=instruct_tag, + quant_type=QuantType.NONE, + is_multimodal=is_multimodal, + ) From 0f0aa0c476c54a41ecf94f3d92690a6038fa9e7d Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 09:27:43 -0700 Subject: [PATCH 15/38] handle quant types per model size --- unsloth/registry/_llama.py | 32 +++++++++++++++++++------------- unsloth/registry/_qwen.py | 6 +++--- unsloth/registry/registry.py | 9 +++++++-- 3 files changed, 29 insertions(+), 18 deletions(-) diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index 211b3ac89e..b62491596c 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -3,6 +3,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register _IS_LLAMA_REGISTERED = False _IS_LLAMA_VISION_REGISTERED = False + class LlamaModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): @@ -27,7 +28,7 @@ LlamaMeta3_1 = ModelMeta( base_name="Llama", instruct_tags=[None, "Instruct"], model_version="3.1", - model_sizes=[8], + model_sizes=["8"], model_info_cls=LlamaModelInfo, is_multimodal=False, quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], @@ -39,10 +40,10 @@ LlamaMeta3_2 = ModelMeta( base_name="Llama", instruct_tags=[None, "Instruct"], model_version="3.2", - model_sizes=[1, 3], + model_sizes=["1", "3"], model_info_cls=LlamaModelInfo, is_multimodal=False, - quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], ) # Llama 3.2 Vision @@ -51,37 +52,42 @@ LlamaMeta3_2_Vision = ModelMeta( base_name="Llama", instruct_tags=[None, "Instruct"], model_version="3.2", - model_sizes=[11, 90], + model_sizes=["11", "90"], model_info_cls=LlamaVisionModelInfo, is_multimodal=True, - quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], + quant_types={ + "11": [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], + "90": [QuantType.NONE], + }, ) -def register_llama_models(): +def register_llama_models(include_original_model: bool = False): global _IS_LLAMA_REGISTERED if _IS_LLAMA_REGISTERED: return - _register_models(LlamaMeta3_1) - _register_models(LlamaMeta3_2) + _register_models(LlamaMeta3_1, include_original_model=include_original_model) + _register_models(LlamaMeta3_2, include_original_model=include_original_model) _IS_LLAMA_REGISTERED = True -def register_llama_vision_models(): +def register_llama_vision_models(include_original_model: bool = False): global _IS_LLAMA_VISION_REGISTERED if _IS_LLAMA_VISION_REGISTERED: return - _register_models(LlamaMeta3_2_Vision) + _register_models(LlamaMeta3_2_Vision, include_original_model=include_original_model) _IS_LLAMA_VISION_REGISTERED = True -register_llama_models() -register_llama_vision_models() + +# register_llama_models(include_original_model=True) +register_llama_vision_models(include_original_model=True) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: print(f"\u2718 {model_id}") else: - print(f"\u2713 {model_id}") \ No newline at end of file + print(f"\u2713 {model_id}") diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py index 2ea340b813..a00d2d5729 100644 --- a/unsloth/registry/_qwen.py +++ b/unsloth/registry/_qwen.py @@ -40,7 +40,7 @@ QwenMeta = ModelMeta( base_name="Qwen", instruct_tags=[None, "Instruct"], model_version="2.5", - model_sizes=[3, 7], + model_sizes=["3", "7"], model_info_cls=QwenModelInfo, is_multimodal=False, quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], @@ -52,7 +52,7 @@ QwenVLMeta = ModelMeta( base_name="Qwen", instruct_tags=["Instruct"], # No base, only instruction tuned model_version="2.5", - model_sizes=[3, 7, 32, 72], + model_sizes=["3", "7", "32", "72"], model_info_cls=QwenVLModelInfo, is_multimodal=True, quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], @@ -64,7 +64,7 @@ QwenQwQMeta = ModelMeta( base_name="QwQ", instruct_tags=[None], model_version="", - model_sizes=[32], + model_sizes=["32"], model_info_cls=QwenQwQModelInfo, is_multimodal=False, quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 3ca7c20f8f..6f50f61d63 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -74,7 +74,7 @@ class ModelMeta: model_info_cls: type[ModelInfo] model_sizes: list[str] = field(default_factory=list) instruct_tags: list[str] = field(default_factory=list) - quant_types: list[QuantType] = field(default_factory=list) + quant_types: list[QuantType] | dict[str, list[QuantType]] = field(default_factory=list) is_multimodal: bool = False @@ -146,7 +146,12 @@ def _register_models(model_meta: ModelMeta, include_original_model: bool = False for size in model_sizes: for instruct_tag in instruct_tags: - for quant_type in quant_types: + # Handle quant types per model size + if isinstance(quant_types, dict): + _quant_types = quant_types[size] + else: + _quant_types = quant_types + for quant_type in _quant_types: _org = "unsloth" # unsloth models -- these are all quantized versions of the original model register_model( model_info_cls=model_info_cls, From 76a2b627661e045bcc75ec40711c14b1929a4bf3 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 09:35:11 -0700 Subject: [PATCH 16/38] separate registration of base and instruct llama3.2 --- unsloth/registry/_llama.py | 25 +++++++++++++++++++------ unsloth/registry/registry.py | 4 ++-- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index b62491596c..6ae838517f 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -34,11 +34,23 @@ LlamaMeta3_1 = ModelMeta( quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], ) -# Llama 3.2 -LlamaMeta3_2 = ModelMeta( +# Llama 3.2 Base Models +LlamaMeta3_2_Base = ModelMeta( org="meta-llama", base_name="Llama", - instruct_tags=[None, "Instruct"], + instruct_tags=[None], + model_version="3.2", + model_sizes=["1", "3"], + model_info_cls=LlamaModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], +) + +# Llama 3.2 Instruction Tuned Models +LlamaMeta3_2_Instruct = ModelMeta( + org="meta-llama", + base_name="Llama", + instruct_tags=["Instruct"], model_version="3.2", model_sizes=["1", "3"], model_info_cls=LlamaModelInfo, @@ -67,7 +79,8 @@ def register_llama_models(include_original_model: bool = False): if _IS_LLAMA_REGISTERED: return _register_models(LlamaMeta3_1, include_original_model=include_original_model) - _register_models(LlamaMeta3_2, include_original_model=include_original_model) + _register_models(LlamaMeta3_2_Base, include_original_model=include_original_model) + _register_models(LlamaMeta3_2_Instruct, include_original_model=include_original_model) _IS_LLAMA_REGISTERED = True @@ -79,8 +92,8 @@ def register_llama_vision_models(include_original_model: bool = False): _IS_LLAMA_VISION_REGISTERED = True -# register_llama_models(include_original_model=True) -register_llama_vision_models(include_original_model=True) +register_llama_models(include_original_model=True) +#register_llama_vision_models(include_original_model=True) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 6f50f61d63..e7a2be0876 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -1,6 +1,6 @@ +import warnings from dataclasses import dataclass, field from enum import Enum -from typing import Literal class QuantType(Enum): @@ -127,7 +127,7 @@ def _check_model_info(model_id: str, properties: list[str] = ["lastModified"]): model_info: HfModelInfo = api.model_info(model_id, expand=properties) except Exception as e: if isinstance(e, RepositoryNotFoundError): - print(f"\u2718 {model_id} not found") + warnings.warn(f"{model_id} not found on Hugging Face") model_info = None else: raise e From 2222e5ad5897b806d5b7ed99019e8e6795a9cf2b Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 09:45:15 -0700 Subject: [PATCH 17/38] add QwenQVQ to registry --- unsloth/registry/_qwen.py | 33 ++++++++++++++++++++++++++++----- unsloth/registry/registry.py | 8 +++++--- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py index a00d2d5729..0b902e3130 100644 --- a/unsloth/registry/_qwen.py +++ b/unsloth/registry/_qwen.py @@ -34,7 +34,17 @@ class QwenQwQModelInfo(ModelInfo): key = cls.append_quant_type(key, quant_type) return key -# Qwen Model Meta +class QwenQVQPreviewModelInfo(ModelInfo): + @classmethod + def construct_model_name( + cls, base_name, version, size, quant_type, instruct_tag + ): + key = f"{base_name}-{size}B-Preview" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + +# Qwen2.5 Model Meta QwenMeta = ModelMeta( org="Qwen", base_name="Qwen", @@ -46,7 +56,7 @@ QwenMeta = ModelMeta( quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], ) -# Qwen VL Model Meta +# Qwen2.5 VL Model Meta QwenVLMeta = ModelMeta( org="Qwen", base_name="Qwen", @@ -70,25 +80,38 @@ QwenQwQMeta = ModelMeta( quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], ) +# Qwen QVQ Preview Model Meta +QwenQVQPreviewMeta = ModelMeta( + org="Qwen", + base_name="QVQ", + instruct_tags=[None], + model_version="", + model_sizes=["72"], + model_info_cls=QwenQVQPreviewModelInfo, + is_multimodal=True, + quant_types=[QuantType.NONE, QuantType.BNB], +) + def register_qwen_models(include_original_model: bool = False): global _IS_QWEN_REGISTERED if _IS_QWEN_REGISTERED: return - _register_models(QwenMeta, include_original_model) + _register_models(QwenMeta, include_original_model=include_original_model) _IS_QWEN_REGISTERED = True def register_qwen_vl_models(include_original_model: bool = False): global _IS_QWEN_VL_REGISTERED if _IS_QWEN_VL_REGISTERED: return - _register_models(QwenVLMeta, include_original_model) + _register_models(QwenVLMeta, include_original_model=include_original_model) _IS_QWEN_VL_REGISTERED = True def register_qwen_qwq_models(include_original_model: bool = False): global _IS_QWEN_QWQ_REGISTERED if _IS_QWEN_QWQ_REGISTERED: return - _register_models(QwenQwQMeta, include_original_model) + _register_models(QwenQwQMeta, include_original_model=include_original_model) + _register_models(QwenQVQPreviewMeta, include_original_model=include_original_model) _IS_QWEN_QWQ_REGISTERED = True # register_qwen_models() diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index e7a2be0876..869a7efb5d 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -5,10 +5,11 @@ from enum import Enum class QuantType(Enum): BNB = "bnb" - UNSLOTH = "unsloth" + UNSLOTH = "unsloth" # dynamic 4-bit quantization GGUF = "GGUF" NONE = "none" +# Tags for Hugging Face model paths BNB_QUANTIZED_TAG = "bnb-4bit" UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG GGUF_TAG = "GGUF" @@ -18,9 +19,9 @@ QUANT_TAG_MAP = { QuantType.UNSLOTH: UNSLOTH_DYNAMIC_QUANT_TAG, QuantType.GGUF: GGUF_TAG, QuantType.NONE: None, -} - +} +# NOTE: models registered with org="unsloth" and QUANT_TYPE.NONE are aliases of QUANT_TYPE.UNSLOTH @dataclass class ModelInfo: org: str @@ -152,6 +153,7 @@ def _register_models(model_meta: ModelMeta, include_original_model: bool = False else: _quant_types = quant_types for quant_type in _quant_types: + # NOTE: models registered with org="unsloth" and QUANT_TYPE.NONE are aliases of QUANT_TYPE.UNSLOTH _org = "unsloth" # unsloth models -- these are all quantized versions of the original model register_model( model_info_cls=model_info_cls, From 756af9f35f0fd8f2da267487d4cca98194400fcd Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 10:10:20 -0700 Subject: [PATCH 18/38] add gemma3 to registry --- unsloth/registry/_gemma.py | 54 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 unsloth/registry/_gemma.py diff --git a/unsloth/registry/_gemma.py b/unsloth/registry/_gemma.py new file mode 100644 index 0000000000..b9abb3737d --- /dev/null +++ b/unsloth/registry/_gemma.py @@ -0,0 +1,54 @@ +from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models + +_IS_GEMMA_REGISTERED = False + +class GemmaModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}-{size}B" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + +# Gemma3 Base Model Meta +GemmaMeta3Base = ModelMeta( + org="google", + base_name="gemma", + instruct_tags=["pt"], # pt = base + model_version="3", + model_sizes=["1", "4", "12", "27"], + model_info_cls=GemmaModelInfo, + is_multimodal=True, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], +) + +# Gemma3 Instruct Model Meta +GemmaMeta3Instruct = ModelMeta( + org="google", + base_name="gemma", + instruct_tags=["it"], # it = instruction tuned + model_version="3", + model_sizes=["1", "4", "12", "27"], + model_info_cls=GemmaModelInfo, + is_multimodal=True, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], +) + +def register_gemma_models(include_original_model: bool = False): + global _IS_GEMMA_REGISTERED + if _IS_GEMMA_REGISTERED: + return + _register_models(GemmaMeta3Base, include_original_model=include_original_model) + _register_models(GemmaMeta3Instruct, include_original_model=include_original_model) + _IS_GEMMA_REGISTERED = True + +register_gemma_models(include_original_model=True) + +if __name__ == "__main__": + from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + for model_id, model_info in MODEL_REGISTRY.items(): + model_info = _check_model_info(model_id) + if model_info is None: + print(f"\u2718 {model_id}") + else: + print(f"\u2713 {model_id}") From a46811c4717ca296ccea0cc9d884c57f54268168 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 10:22:50 -0700 Subject: [PATCH 19/38] add phi --- unsloth/registry/_phi.py | 62 ++++++++++++++++++++++++++++++ unsloth/registry/model_registry.py | 59 ++-------------------------- 2 files changed, 66 insertions(+), 55 deletions(-) create mode 100644 unsloth/registry/_phi.py diff --git a/unsloth/registry/_phi.py b/unsloth/registry/_phi.py new file mode 100644 index 0000000000..a6d18cbd61 --- /dev/null +++ b/unsloth/registry/_phi.py @@ -0,0 +1,62 @@ +from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models + +_IS_PHI_REGISTERED = False +_IS_PHI_INSTRUCT_REGISTERED = False + +class PhiModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + +# Phi Model Meta +PhiMeta = ModelMeta( + org="microsoft", + base_name="phi", + instruct_tags=[None], + model_version="4", + model_sizes=["1"], # Assuming only one size + model_info_cls=PhiModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], +) + +# Phi Instruct Model Meta +PhiInstructMeta = ModelMeta( + org="microsoft", + base_name="phi", + instruct_tags=["mini-instruct"], + model_version="4", + model_sizes=["1"], # Assuming only one size + model_info_cls=PhiModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], +) + +def register_phi_models(include_original_model: bool = False): + global _IS_PHI_REGISTERED + if _IS_PHI_REGISTERED: + return + _register_models(PhiMeta, include_original_model=include_original_model) + _IS_PHI_REGISTERED = True + +def register_phi_instruct_models(include_original_model: bool = False): + global _IS_PHI_INSTRUCT_REGISTERED + if _IS_PHI_INSTRUCT_REGISTERED: + return + _register_models(PhiInstructMeta, include_original_model=include_original_model) + _IS_PHI_INSTRUCT_REGISTERED = True + +register_phi_models(include_original_model=True) +register_phi_instruct_models(include_original_model=True) + +if __name__ == "__main__": + from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + for model_id, model_info in MODEL_REGISTRY.items(): + model_info = _check_model_info(model_id) + if model_info is None: + print(f"\u2718 {model_id}") + else: + print(f"\u2713 {model_id}") \ No newline at end of file diff --git a/unsloth/registry/model_registry.py b/unsloth/registry/model_registry.py index a0cd71c17a..de9609934c 100644 --- a/unsloth/registry/model_registry.py +++ b/unsloth/registry/model_registry.py @@ -4,11 +4,11 @@ from typing import Callable, Literal from unsloth.registry._llama import LlamaMeta3_1, LlamaMeta3_2 from unsloth.registry.common import ModelInfo, ModelMeta -_IS_LLAMA_REGISTERED = False -_IS_LLAMA_VISION_REGISTERED = False +# _IS_LLAMA_REGISTERED = False +# _IS_LLAMA_VISION_REGISTERED = False -_IS_QWEN_REGISTERED = False -_IS_QWEN_VL_REGISTERED = False +# _IS_QWEN_REGISTERED = False +# _IS_QWEN_VL_REGISTERED = False _IS_GEMMA_REGISTERED = False @@ -17,28 +17,6 @@ _IS_PHI_INSTRUCT_REGISTERED = False -# class QwenModelInfo(ModelInfo): -# @classmethod -# def construct_model_name( -# cls, base_name, version, size, quant_type, instruct_tag -# ): -# key = f"{base_name}{version}-{size}B" -# key = cls.append_instruct_tag(key, instruct_tag) -# key = cls.append_quant_type(key, quant_type) -# return key - - -# class QwenVLModelInfo(ModelInfo): -# @classmethod -# def construct_model_name( -# cls, base_name, version, size, quant_type, instruct_tag -# ): -# key = f"{base_name}{version}-VL-{size}B" -# key = cls.append_instruct_tag(key, instruct_tag) -# key = cls.append_quant_type(key, quant_type) -# return key - - # class PhiModelInfo(ModelInfo): # @classmethod # def construct_model_name( @@ -55,35 +33,6 @@ _IS_PHI_INSTRUCT_REGISTERED = False # # Qwen text only models # # NOTE: Qwen vision models will be registered separately -# _QWEN_INFO = { -# "org": "Qwen", -# "base_name": "Qwen", -# "instruct_tags": [None, "Instruct"], -# "model_versions": ["2.5"], -# "model_sizes": {"2.5": [3, 7]}, -# "is_multimodal": False, -# "model_info_cls": QwenModelInfo, -# } - -# _QWEN_VL_INFO = { -# "org": "Qwen", -# "base_name": "Qwen", -# "instruct_tags": ["Instruct"], # No base, only instruction tuned -# "model_versions": ["2.5"], -# "model_sizes": {"2.5": [3, 7, 32, 72]}, -# "is_multimodal": True, -# "instruction_tuned_only": True, -# "model_info_cls": QwenVLModelInfo, -# } - -# _GEMMA_INFO = { -# "org": "google", -# "base_name": "gemma", -# "instruct_tags": ["pt", "it"], # pt = base, it = instruction tuned -# "model_versions": ["3"], -# "model_sizes": {"3": [1, 4, 12, 27]}, -# "is_multimodal": True, -# } # _PHI_INFO = { # "org": "microsoft", From e2ff538fc5a25fb2adabdc679d65c8ce2a22a75e Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 11:23:22 -0700 Subject: [PATCH 20/38] add deepseek v3 --- unsloth/registry/_deepseek.py | 53 +++++++++++++++++++++++++++++++++++ unsloth/registry/registry.py | 3 ++ 2 files changed, 56 insertions(+) create mode 100644 unsloth/registry/_deepseek.py diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py new file mode 100644 index 0000000000..8bdcd3c2e7 --- /dev/null +++ b/unsloth/registry/_deepseek.py @@ -0,0 +1,53 @@ +from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models + +_IS_DEEPSEEKV3_REGISTERED = False + +class DeepseekV3ModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-V{version}" + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + +# Deepseek V3 Model Meta +DeepseekV3Meta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek", + instruct_tags=[None], + model_version="3", + model_sizes=[""], + model_info_cls=DeepseekV3ModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BF16], +) + +DeepseekV3_0324Meta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek", + instruct_tags=[None], + model_version="3-0324", + model_sizes=[""], + model_info_cls=DeepseekV3ModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.GGUF], +) + +def register_deepseek_v3_models(include_original_model: bool = False): + global _IS_DEEPSEEKV3_REGISTERED + if _IS_DEEPSEEKV3_REGISTERED: + return + _register_models(DeepseekV3Meta, include_original_model=include_original_model) + _register_models(DeepseekV3_0324Meta, include_original_model=include_original_model) + _IS_DEEPSEEKV3_REGISTERED = True + +register_deepseek_v3_models(include_original_model=True) + +if __name__ == "__main__": + from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + for model_id, model_info in MODEL_REGISTRY.items(): + model_info = _check_model_info(model_id) + if model_info is None: + print(f"\u2718 {model_id}") + else: + print(f"\u2713 {model_id}") diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 869a7efb5d..1eee884259 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -8,17 +8,20 @@ class QuantType(Enum): UNSLOTH = "unsloth" # dynamic 4-bit quantization GGUF = "GGUF" NONE = "none" + BF16 = "bf16" # only for Deepseek V3 # Tags for Hugging Face model paths BNB_QUANTIZED_TAG = "bnb-4bit" UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG GGUF_TAG = "GGUF" +BF16_TAG = "bf16" QUANT_TAG_MAP = { QuantType.BNB: BNB_QUANTIZED_TAG, QuantType.UNSLOTH: UNSLOTH_DYNAMIC_QUANT_TAG, QuantType.GGUF: GGUF_TAG, QuantType.NONE: None, + QuantType.BF16: BF16_TAG, } # NOTE: models registered with org="unsloth" and QUANT_TYPE.NONE are aliases of QUANT_TYPE.UNSLOTH From 9f8f78c90b47e3d00bcd86df5c3a95d463379277 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 11:30:47 -0700 Subject: [PATCH 21/38] add deepseek r1 base --- unsloth/registry/_deepseek.py | 32 +++++++++++++++++++++++++++++++- unsloth/registry/registry.py | 3 ++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 8bdcd3c2e7..0346520060 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -1,6 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models _IS_DEEPSEEKV3_REGISTERED = False +_IS_DEEPSEEKR1_REGISTERED = False class DeepseekV3ModelInfo(ModelInfo): @classmethod @@ -10,6 +11,14 @@ class DeepseekV3ModelInfo(ModelInfo): key = cls.append_quant_type(key, quant_type) return key +class DeepseekR1ModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}" if version else base_name + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + # Deepseek V3 Model Meta DeepseekV3Meta = ModelMeta( org="deepseek-ai", @@ -33,6 +42,17 @@ DeepseekV3_0324Meta = ModelMeta( quant_types=[QuantType.NONE, QuantType.GGUF], ) +DeepseekR1Meta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek-R1", + instruct_tags=[None], + model_version="", + model_sizes=[""], + model_info_cls=DeepseekR1ModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BF16, QuantType.GGUF], +) + def register_deepseek_v3_models(include_original_model: bool = False): global _IS_DEEPSEEKV3_REGISTERED if _IS_DEEPSEEKV3_REGISTERED: @@ -41,7 +61,17 @@ def register_deepseek_v3_models(include_original_model: bool = False): _register_models(DeepseekV3_0324Meta, include_original_model=include_original_model) _IS_DEEPSEEKV3_REGISTERED = True -register_deepseek_v3_models(include_original_model=True) + +def register_deepseek_r1_models(include_original_model: bool = False): + global _IS_DEEPSEEKR1_REGISTERED + if _IS_DEEPSEEKR1_REGISTERED: + return + _register_models(DeepseekR1Meta, include_original_model=include_original_model) + _IS_DEEPSEEKR1_REGISTERED = True + +#register_deepseek_v3_models(include_original_model=True) +register_deepseek_r1_models(include_original_model=True) + if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 1eee884259..1e2c667e13 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -35,7 +35,8 @@ class ModelInfo: is_multimodal: bool = False instruct_tag: str = None quant_type: QuantType = None - + description: str = None + def __post_init__(self): self.name = self.name or self.construct_model_name( self.base_name, From 767044e7f277dd463bc37b7357eeb12944f5c333 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 11:32:21 -0700 Subject: [PATCH 22/38] add deepseek r1 zero --- unsloth/registry/_deepseek.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 0346520060..bd0ea31cbf 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -53,6 +53,16 @@ DeepseekR1Meta = ModelMeta( quant_types=[QuantType.NONE, QuantType.BF16, QuantType.GGUF], ) +DeepseekR1ZeroMeta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek-R1", + instruct_tags=[None], + model_version="Zero", + model_sizes=[""], + model_info_cls=DeepseekR1ModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.GGUF], +) def register_deepseek_v3_models(include_original_model: bool = False): global _IS_DEEPSEEKV3_REGISTERED if _IS_DEEPSEEKV3_REGISTERED: @@ -67,6 +77,7 @@ def register_deepseek_r1_models(include_original_model: bool = False): if _IS_DEEPSEEKR1_REGISTERED: return _register_models(DeepseekR1Meta, include_original_model=include_original_model) + _register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model) _IS_DEEPSEEKR1_REGISTERED = True #register_deepseek_v3_models(include_original_model=True) From 7dec39e3b63bf57d70a83c587d0b11b9c68ede78 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 11:47:51 -0700 Subject: [PATCH 23/38] add deepseek distill llama --- unsloth/registry/_deepseek.py | 38 ++++++++++++++++++++++++++++++++++- unsloth/utils/hf_hub.py | 24 +++++++++++----------- 2 files changed, 49 insertions(+), 13 deletions(-) diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index bd0ea31cbf..b3bf398cf1 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -2,7 +2,8 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register _IS_DEEPSEEKV3_REGISTERED = False _IS_DEEPSEEKR1_REGISTERED = False - +_IS_DEEPSEEKR1_ZERO_REGISTERED = False +_IS_DEEPSEEKR1_DISTILL_REGISTERED = False class DeepseekV3ModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): @@ -15,6 +16,8 @@ class DeepseekR1ModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-{version}" if version else base_name + if size: + key = f"{key}-{size}B" key = cls.append_instruct_tag(key, instruct_tag) key = cls.append_quant_type(key, quant_type) return key @@ -63,6 +66,28 @@ DeepseekR1ZeroMeta = ModelMeta( is_multimodal=False, quant_types=[QuantType.NONE, QuantType.GGUF], ) + +DeepseekR1DistillMeta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek-R1-Distill", + instruct_tags=[None], + model_version="Llama", + model_sizes=["8", "70"], + model_info_cls=DeepseekR1ModelInfo, + is_multimodal=False, + quant_types={"8": [QuantType.UNSLOTH, QuantType.GGUF], "70": [QuantType.GGUF]}, +) + + # "Qwen-7B-unsloth-bnb-4bit", + # "Qwen-1.5B-unsloth-bnb-4bit", + # "Qwen-32B-GGUF", + # "Llama-8B-GGUF", + # "Qwen-14B-GGUF", + # "Qwen-32B-bnb-4bit", + # "Qwen-1.5B-GGUF", + # "Qwen-14B-unsloth-bnb-4bit", + # "Llama-70B-GGUF" + def register_deepseek_v3_models(include_original_model: bool = False): global _IS_DEEPSEEKV3_REGISTERED if _IS_DEEPSEEKV3_REGISTERED: @@ -78,11 +103,22 @@ def register_deepseek_r1_models(include_original_model: bool = False): return _register_models(DeepseekR1Meta, include_original_model=include_original_model) _register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model) + _register_models(DeepseekR1DistillMeta, include_original_model=include_original_model) _IS_DEEPSEEKR1_REGISTERED = True #register_deepseek_v3_models(include_original_model=True) register_deepseek_r1_models(include_original_model=True) +def _list_deepseek_r1_distill_models(): + from unsloth.utils.hf_hub import ModelInfo as HfModelInfo + from unsloth.utils.hf_hub import list_models + models: list[HfModelInfo] = list_models(author="unsloth", search="Distill") + for model in models: + model_id = model.id + model_name = model_id.split("/")[-1] + # parse out only the version + version = model_name.removeprefix("DeepSeek-R1-Distill-") + print(version) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info diff --git a/unsloth/utils/hf_hub.py b/unsloth/utils/hf_hub.py index da3f72a18e..30255b8636 100644 --- a/unsloth/utils/hf_hub.py +++ b/unsloth/utils/hf_hub.py @@ -1,6 +1,6 @@ from huggingface_hub import HfApi, ModelInfo -api: HfApi +_HFAPI: HfApi = None POPULARITY_PROPERTIES = [ "downloads", @@ -32,27 +32,27 @@ def get_model_info( Default properties: ["safetensors", "lastModified"], only retrieves minimal information. Set to None to retrieve the full model information. """ - global api - if api is None: - api = HfApi() + global _HFAPI + if _HFAPI is None: + _HFAPI = HfApi() try: - model_info: ModelInfo = api.model_info(model_id, expand=properties) + model_info: ModelInfo = _HFAPI.model_info(model_id, expand=properties) except Exception as e: print(f"Error getting model info for {model_id}: {e}") model_info = None return model_info -def retrieve_models( +def list_models( properties: list[str] = None, full: bool = False, sort: str = "downloads", author: str = "unsloth", search: str = None, limit: int = 10, -) -> ModelInfo: +) -> list[ModelInfo]: """ - Retrieve models from the Hugging Face Hub. + Retrieve model information from the Hugging Face Hub. properties: list[str] = See https://huggingface.co/docs/huggingface_hub/api-ref/hf_hub/hf_api/list_models full: bool = Whether to retrieve the full model information, if True properties will be ignored. @@ -61,13 +61,13 @@ def retrieve_models( search: str = The search query for filtering models. """ - global api - if api is None: - api = HfApi() + global _HFAPI + if _HFAPI is None: + _HFAPI = HfApi() if full: properties = None - models: list[ModelInfo] = api.list_models( + models: list[ModelInfo] = _HFAPI.list_models( author=author, search=search, sort=sort, From 7157f3c47cf7965f8d8eccc22205e2a0ed9d224f Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 12:04:57 -0700 Subject: [PATCH 24/38] add deepseek distill models --- unsloth/registry/_deepseek.py | 64 +++++++++++++++++++++++++++++------ 1 file changed, 54 insertions(+), 10 deletions(-) diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index b3bf398cf1..8e87ba11dd 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -3,7 +3,9 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register _IS_DEEPSEEKV3_REGISTERED = False _IS_DEEPSEEKR1_REGISTERED = False _IS_DEEPSEEKR1_ZERO_REGISTERED = False -_IS_DEEPSEEKR1_DISTILL_REGISTERED = False +_IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED = False +_IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED = False + class DeepseekV3ModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): @@ -67,7 +69,7 @@ DeepseekR1ZeroMeta = ModelMeta( quant_types=[QuantType.NONE, QuantType.GGUF], ) -DeepseekR1DistillMeta = ModelMeta( +DeepseekR1DistillLlamaMeta = ModelMeta( org="deepseek-ai", base_name="DeepSeek-R1-Distill", instruct_tags=[None], @@ -78,16 +80,27 @@ DeepseekR1DistillMeta = ModelMeta( quant_types={"8": [QuantType.UNSLOTH, QuantType.GGUF], "70": [QuantType.GGUF]}, ) +# Deepseek R1 Distill Qwen Model Meta +DeepseekR1DistillQwenMeta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek-R1-Distill", + instruct_tags=[None], + model_version="Qwen", + model_sizes=["1.5", "7", "14", "32"], + model_info_cls=DeepseekR1ModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF] +) + # "Qwen-7B-unsloth-bnb-4bit", # "Qwen-1.5B-unsloth-bnb-4bit", # "Qwen-32B-GGUF", - # "Llama-8B-GGUF", + # "Qwen-14B-GGUF", # "Qwen-32B-bnb-4bit", # "Qwen-1.5B-GGUF", # "Qwen-14B-unsloth-bnb-4bit", - # "Llama-70B-GGUF" - + def register_deepseek_v3_models(include_original_model: bool = False): global _IS_DEEPSEEKV3_REGISTERED if _IS_DEEPSEEKV3_REGISTERED: @@ -102,23 +115,50 @@ def register_deepseek_r1_models(include_original_model: bool = False): if _IS_DEEPSEEKR1_REGISTERED: return _register_models(DeepseekR1Meta, include_original_model=include_original_model) - _register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model) - _register_models(DeepseekR1DistillMeta, include_original_model=include_original_model) _IS_DEEPSEEKR1_REGISTERED = True -#register_deepseek_v3_models(include_original_model=True) +def register_deepseek_r1_zero_models(include_original_model: bool = False): + global _IS_DEEPSEEKR1_ZERO_REGISTERED + if _IS_DEEPSEEKR1_ZERO_REGISTERED: + return + _register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model) + _IS_DEEPSEEKR1_ZERO_REGISTERED = True + +def register_deepseek_r1_distill_llama_models(include_original_model: bool = False): + global _IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED + if _IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED: + return + _register_models(DeepseekR1DistillLlamaMeta, include_original_model=include_original_model) + _IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED = True + +def register_deepseek_r1_distill_qwen_models(include_original_model: bool = False): + global _IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED + if _IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED: + return + _register_models(DeepseekR1DistillQwenMeta, include_original_model=include_original_model) + _IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED = True + +def register_deepseek_r1_distill_models(include_original_model: bool = False): + register_deepseek_r1_distill_qwen_models(include_original_model=include_original_model) + register_deepseek_r1_distill_llama_models(include_original_model=include_original_model) + +register_deepseek_v3_models(include_original_model=True) register_deepseek_r1_models(include_original_model=True) +register_deepseek_r1_distill_models(include_original_model=True) def _list_deepseek_r1_distill_models(): from unsloth.utils.hf_hub import ModelInfo as HfModelInfo from unsloth.utils.hf_hub import list_models - models: list[HfModelInfo] = list_models(author="unsloth", search="Distill") + models: list[HfModelInfo] = list_models(author="unsloth", search="Distill", limit=1000) + distill_models = [] for model in models: model_id = model.id model_name = model_id.split("/")[-1] # parse out only the version version = model_name.removeprefix("DeepSeek-R1-Distill-") - print(version) + distill_models.append(version) + + return distill_models if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info @@ -128,3 +168,7 @@ if __name__ == "__main__": print(f"\u2718 {model_id}") else: print(f"\u2713 {model_id}") + # distill_models = _list_deepseek_r1_distill_models() + # for model in sorted(distill_models): + # if "qwen" in model.lower(): + # print(model) \ No newline at end of file From 025e22b666221fa225d3e178319c0d3dbc080dda Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 15:06:08 -0700 Subject: [PATCH 25/38] remove redundant code when constructing model names --- unsloth/registry/_deepseek.py | 8 ++------ unsloth/registry/_gemma.py | 4 +--- unsloth/registry/_llama.py | 8 ++------ unsloth/registry/_phi.py | 4 +--- unsloth/registry/_qwen.py | 32 ++++++++------------------------ unsloth/registry/registry.py | 8 +++++--- 6 files changed, 19 insertions(+), 45 deletions(-) diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 8e87ba11dd..148093155c 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -10,9 +10,7 @@ class DeepseekV3ModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-V{version}" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) class DeepseekR1ModelInfo(ModelInfo): @classmethod @@ -20,9 +18,7 @@ class DeepseekR1ModelInfo(ModelInfo): key = f"{base_name}-{version}" if version else base_name if size: key = f"{key}-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) # Deepseek V3 Model Meta DeepseekV3Meta = ModelMeta( diff --git a/unsloth/registry/_gemma.py b/unsloth/registry/_gemma.py index b9abb3737d..4fef26d533 100644 --- a/unsloth/registry/_gemma.py +++ b/unsloth/registry/_gemma.py @@ -6,9 +6,7 @@ class GemmaModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-{version}-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) # Gemma3 Base Model Meta GemmaMeta3Base = ModelMeta( diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index 6ae838517f..dbf7c8a9d6 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -8,18 +8,14 @@ class LlamaModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-{version}-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) class LlamaVisionModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-{version}-{size}B-Vision" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) # Llama 3.1 diff --git a/unsloth/registry/_phi.py b/unsloth/registry/_phi.py index a6d18cbd61..c69eaf83bb 100644 --- a/unsloth/registry/_phi.py +++ b/unsloth/registry/_phi.py @@ -7,9 +7,7 @@ class PhiModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-{version}" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) # Phi Model Meta PhiMeta = ModelMeta( diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py index 0b902e3130..c9a0a4d4ec 100644 --- a/unsloth/registry/_qwen.py +++ b/unsloth/registry/_qwen.py @@ -5,44 +5,28 @@ _IS_QWEN_VL_REGISTERED = False _IS_QWEN_QWQ_REGISTERED = False class QwenModelInfo(ModelInfo): @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}{version}-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) class QwenVLModelInfo(ModelInfo): @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}{version}-VL-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) class QwenQwQModelInfo(ModelInfo): @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-{size}B" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) class QwenQVQPreviewModelInfo(ModelInfo): @classmethod - def construct_model_name( - cls, base_name, version, size, quant_type, instruct_tag - ): + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): key = f"{base_name}-{size}B-Preview" - key = cls.append_instruct_tag(key, instruct_tag) - key = cls.append_quant_type(key, quant_type) - return key + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) # Qwen2.5 Model Meta QwenMeta = ModelMeta( diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 1e2c667e13..590beebeeb 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -36,7 +36,7 @@ class ModelInfo: instruct_tag: str = None quant_type: QuantType = None description: str = None - + def __post_init__(self): self.name = self.name or self.construct_model_name( self.base_name, @@ -61,8 +61,10 @@ class ModelInfo: return key @classmethod - def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): - raise NotImplementedError("Subclass must implement this method") + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag, key=""): + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key @property def model_path( From 5c402c9e82b0464c2fff0253990925bbdbbb53a8 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 15:31:01 -0700 Subject: [PATCH 26/38] add mistral small to registry --- unsloth/registry/_deepseek.py | 9 ++--- unsloth/registry/_mistral.py | 66 +++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) create mode 100644 unsloth/registry/_mistral.py diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 148093155c..35cbc17484 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -138,10 +138,6 @@ def register_deepseek_r1_distill_models(include_original_model: bool = False): register_deepseek_r1_distill_qwen_models(include_original_model=include_original_model) register_deepseek_r1_distill_llama_models(include_original_model=include_original_model) -register_deepseek_v3_models(include_original_model=True) -register_deepseek_r1_models(include_original_model=True) -register_deepseek_r1_distill_models(include_original_model=True) - def _list_deepseek_r1_distill_models(): from unsloth.utils.hf_hub import ModelInfo as HfModelInfo from unsloth.utils.hf_hub import list_models @@ -156,6 +152,11 @@ def _list_deepseek_r1_distill_models(): return distill_models + +register_deepseek_v3_models(include_original_model=True) +register_deepseek_r1_models(include_original_model=True) +register_deepseek_r1_distill_models(include_original_model=True) + if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info for model_id, model_info in MODEL_REGISTRY.items(): diff --git a/unsloth/registry/_mistral.py b/unsloth/registry/_mistral.py new file mode 100644 index 0000000000..65f1256708 --- /dev/null +++ b/unsloth/registry/_mistral.py @@ -0,0 +1,66 @@ +import copy + +from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models + +_IS_MISTRAL_SMALL_REGISTERED = False + +_MISTRAL_SMALL_03_25_VERSION = "2503" +_MISTRAL_SMALL_01_25_VERSION = "2501" +_MISTRAL_SMALL_09_24_VERSION = "2409" # Not uploaded to unsloth + +class MistralSmallModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + if version == _MISTRAL_SMALL_03_25_VERSION: + key = f"{base_name}-3.1-{size}B-{instruct_tag}" + else: + key = f"{base_name}-{size}B-{instruct_tag}" + key += f"-{version}" + key = cls.append_quant_type(key, quant_type) + + return key + + +MistralSmall_2503_Base_Meta = ModelMeta( + org="mistralai", + base_name="Mistral-Small", + instruct_tags=["Base"], + model_version=_MISTRAL_SMALL_03_25_VERSION, + model_sizes=["24"], + model_info_cls=MistralSmallModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB], +) + +MistralSmall_2503_Instruct_Meta = copy.deepcopy(MistralSmall_2503_Base_Meta) +MistralSmall_2503_Instruct_Meta.instruct_tags = ["Instruct"] +MistralSmall_2503_Instruct_Meta.quant_types = [QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF] + +MistralSmall_2501_Base_Meta = copy.deepcopy(MistralSmall_2503_Base_Meta) +MistralSmall_2501_Base_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION + +MistralSmall_2501_Instruct_Meta = copy.deepcopy(MistralSmall_2503_Instruct_Meta) +MistralSmall_2501_Instruct_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION + +def register_mistral_small_models(): + global _IS_MISTRAL_SMALL_REGISTERED + if _IS_MISTRAL_SMALL_REGISTERED: + return + _register_models(MistralSmall_2503_Base_Meta) + _register_models(MistralSmall_2503_Instruct_Meta) + _register_models(MistralSmall_2501_Base_Meta) + _register_models(MistralSmall_2501_Instruct_Meta) + + _IS_MISTRAL_SMALL_REGISTERED = True + +register_mistral_small_models() + + +if __name__ == "__main__": + from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + for model_id, model_info in MODEL_REGISTRY.items(): + model_info = _check_model_info(model_id) + if model_info is None: + print(f"\u2718 {model_id}") + else: + print(f"\u2713 {model_id}") \ No newline at end of file From 16f644e95db5a14a9b60719c636b62427cb8a415 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:01:51 -0700 Subject: [PATCH 27/38] rename model registration methods --- unsloth/registry/__init__.py | 5 ++++ unsloth/registry/_gemma.py | 25 +++++++++++++----- unsloth/registry/_llama.py | 51 ++++++++++++++++++------------------ unsloth/registry/_qwen.py | 36 +++++++++++++------------ 4 files changed, 68 insertions(+), 49 deletions(-) diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py index e69de29bb2..dd5b45c4ee 100644 --- a/unsloth/registry/__init__.py +++ b/unsloth/registry/__init__.py @@ -0,0 +1,5 @@ +# from ._deepseek import register_deepseek_models, register +# from ._llama import register_llama_models, register_llama_vision_models +# from ._mistral import register_mistral_models +# from ._openai import register_openai_models +# from ._qwen import register_qwen_models diff --git a/unsloth/registry/_gemma.py b/unsloth/registry/_gemma.py index 4fef26d533..8c47e7e69d 100644 --- a/unsloth/registry/_gemma.py +++ b/unsloth/registry/_gemma.py @@ -1,6 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models -_IS_GEMMA_REGISTERED = False +_IS_GEMMA_3_BASE_REGISTERED = False +_IS_GEMMA_3_INSTRUCT_REGISTERED = False class GemmaModelInfo(ModelInfo): @classmethod @@ -32,17 +33,27 @@ GemmaMeta3Instruct = ModelMeta( quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], ) -def register_gemma_models(include_original_model: bool = False): - global _IS_GEMMA_REGISTERED - if _IS_GEMMA_REGISTERED: +def register_gemma_3_base_models(include_original_model: bool = False): + global _IS_GEMMA_3_BASE_REGISTERED + if _IS_GEMMA_3_BASE_REGISTERED: return _register_models(GemmaMeta3Base, include_original_model=include_original_model) - _register_models(GemmaMeta3Instruct, include_original_model=include_original_model) - _IS_GEMMA_REGISTERED = True + _IS_GEMMA_3_BASE_REGISTERED = True + +def register_gemma_3_instruct_models(include_original_model: bool = False): + global _IS_GEMMA_3_INSTRUCT_REGISTERED + if _IS_GEMMA_3_INSTRUCT_REGISTERED: + return + _register_models(GemmaMeta3Instruct, include_original_model=include_original_model) + _IS_GEMMA_3_INSTRUCT_REGISTERED = True + +def register_gemma_models(include_original_model: bool = False): + register_gemma_3_base_models(include_original_model=include_original_model) + register_gemma_3_instruct_models(include_original_model=include_original_model) -register_gemma_models(include_original_model=True) if __name__ == "__main__": + register_gemma_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index dbf7c8a9d6..c84c5b8d30 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -1,7 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models -_IS_LLAMA_REGISTERED = False -_IS_LLAMA_VISION_REGISTERED = False +_IS_LLAMA_3_REGISTERED = False +_IS_LLAMA_3_2_VISION_REGISTERED = False class LlamaModelInfo(ModelInfo): @@ -19,7 +19,7 @@ class LlamaVisionModelInfo(ModelInfo): # Llama 3.1 -LlamaMeta3_1 = ModelMeta( +LlamaMeta_3_1 = ModelMeta( org="meta-llama", base_name="Llama", instruct_tags=[None, "Instruct"], @@ -31,7 +31,7 @@ LlamaMeta3_1 = ModelMeta( ) # Llama 3.2 Base Models -LlamaMeta3_2_Base = ModelMeta( +LlamaMeta_3_2_Base = ModelMeta( org="meta-llama", base_name="Llama", instruct_tags=[None], @@ -43,7 +43,7 @@ LlamaMeta3_2_Base = ModelMeta( ) # Llama 3.2 Instruction Tuned Models -LlamaMeta3_2_Instruct = ModelMeta( +LlamaMeta_3_2_Instruct = ModelMeta( org="meta-llama", base_name="Llama", instruct_tags=["Instruct"], @@ -55,7 +55,7 @@ LlamaMeta3_2_Instruct = ModelMeta( ) # Llama 3.2 Vision -LlamaMeta3_2_Vision = ModelMeta( +LlamaMeta_3_2_Vision = ModelMeta( org="meta-llama", base_name="Llama", instruct_tags=[None, "Instruct"], @@ -70,28 +70,29 @@ LlamaMeta3_2_Vision = ModelMeta( ) +def register_llama_3_models(include_original_model: bool = False): + global _IS_LLAMA_3_REGISTERED + if _IS_LLAMA_3_REGISTERED: + return + _register_models(LlamaMeta_3_1, include_original_model=include_original_model) + _register_models(LlamaMeta_3_2_Base, include_original_model=include_original_model) + _register_models(LlamaMeta_3_2_Instruct, include_original_model=include_original_model) + _IS_LLAMA_3_REGISTERED = True + +def register_llama_3_2_vision_models(include_original_model: bool = False): + global _IS_LLAMA_3_2_VISION_REGISTERED + if _IS_LLAMA_3_2_VISION_REGISTERED: + return + _register_models(LlamaMeta_3_2_Vision, include_original_model=include_original_model) + _IS_LLAMA_3_2_VISION_REGISTERED = True + + def register_llama_models(include_original_model: bool = False): - global _IS_LLAMA_REGISTERED - if _IS_LLAMA_REGISTERED: - return - _register_models(LlamaMeta3_1, include_original_model=include_original_model) - _register_models(LlamaMeta3_2_Base, include_original_model=include_original_model) - _register_models(LlamaMeta3_2_Instruct, include_original_model=include_original_model) - _IS_LLAMA_REGISTERED = True - - -def register_llama_vision_models(include_original_model: bool = False): - global _IS_LLAMA_VISION_REGISTERED - if _IS_LLAMA_VISION_REGISTERED: - return - _register_models(LlamaMeta3_2_Vision, include_original_model=include_original_model) - _IS_LLAMA_VISION_REGISTERED = True - - -register_llama_models(include_original_model=True) -#register_llama_vision_models(include_original_model=True) + register_llama_3_models(include_original_model=include_original_model) + register_llama_3_2_vision_models(include_original_model=include_original_model) if __name__ == "__main__": + register_llama_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info for model_id, model_info in MODEL_REGISTRY.items(): diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py index c9a0a4d4ec..c364f9b099 100644 --- a/unsloth/registry/_qwen.py +++ b/unsloth/registry/_qwen.py @@ -1,7 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models -_IS_QWEN_REGISTERED = False -_IS_QWEN_VL_REGISTERED = False +_IS_QWEN_2_5_REGISTERED = False +_IS_QWEN_2_5_VL_REGISTERED = False _IS_QWEN_QWQ_REGISTERED = False class QwenModelInfo(ModelInfo): @classmethod @@ -29,7 +29,7 @@ class QwenQVQPreviewModelInfo(ModelInfo): return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) # Qwen2.5 Model Meta -QwenMeta = ModelMeta( +Qwen_2_5_Meta = ModelMeta( org="Qwen", base_name="Qwen", instruct_tags=[None, "Instruct"], @@ -41,7 +41,7 @@ QwenMeta = ModelMeta( ) # Qwen2.5 VL Model Meta -QwenVLMeta = ModelMeta( +Qwen_2_5_VLMeta = ModelMeta( org="Qwen", base_name="Qwen", instruct_tags=["Instruct"], # No base, only instruction tuned @@ -76,19 +76,19 @@ QwenQVQPreviewMeta = ModelMeta( quant_types=[QuantType.NONE, QuantType.BNB], ) -def register_qwen_models(include_original_model: bool = False): - global _IS_QWEN_REGISTERED - if _IS_QWEN_REGISTERED: +def register_qwen_2_5_models(include_original_model: bool = False): + global _IS_QWEN_2_5_REGISTERED + if _IS_QWEN_2_5_REGISTERED: return - _register_models(QwenMeta, include_original_model=include_original_model) - _IS_QWEN_REGISTERED = True + _register_models(Qwen_2_5_Meta, include_original_model=include_original_model) + _IS_QWEN_2_5_REGISTERED = True -def register_qwen_vl_models(include_original_model: bool = False): - global _IS_QWEN_VL_REGISTERED - if _IS_QWEN_VL_REGISTERED: +def register_qwen_2_5_vl_models(include_original_model: bool = False): + global _IS_QWEN_2_5_VL_REGISTERED + if _IS_QWEN_2_5_VL_REGISTERED: return - _register_models(QwenVLMeta, include_original_model=include_original_model) - _IS_QWEN_VL_REGISTERED = True + _register_models(Qwen_2_5_VLMeta, include_original_model=include_original_model) + _IS_QWEN_2_5_VL_REGISTERED = True def register_qwen_qwq_models(include_original_model: bool = False): global _IS_QWEN_QWQ_REGISTERED @@ -98,11 +98,13 @@ def register_qwen_qwq_models(include_original_model: bool = False): _register_models(QwenQVQPreviewMeta, include_original_model=include_original_model) _IS_QWEN_QWQ_REGISTERED = True -# register_qwen_models() -# register_qwen_vl_models() -register_qwen_qwq_models(include_original_model=True) +def register_qwen_models(include_original_model: bool = False): + register_qwen_2_5_models(include_original_model=include_original_model) + register_qwen_2_5_vl_models(include_original_model=include_original_model) + register_qwen_qwq_models(include_original_model=include_original_model) if __name__ == "__main__": + register_qwen_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) From 9a276978d219b42a1065630158646f84f01864a7 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:03:05 -0700 Subject: [PATCH 28/38] rename deepseek registration methods --- unsloth/registry/_deepseek.py | 67 +++++++++++++++++++++-------------- 1 file changed, 40 insertions(+), 27 deletions(-) diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 35cbc17484..1f97a02f1c 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -1,10 +1,11 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models -_IS_DEEPSEEKV3_REGISTERED = False -_IS_DEEPSEEKR1_REGISTERED = False -_IS_DEEPSEEKR1_ZERO_REGISTERED = False -_IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED = False -_IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED = False +_IS_DEEPSEEK_V3_REGISTERED = False +_IS_DEEPSEEK_V3_0324_REGISTERED = False +_IS_DEEPSEEK_R1_REGISTERED = False +_IS_DEEPSEEK_R1_ZERO_REGISTERED = False +_IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED = False +_IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED = False class DeepseekV3ModelInfo(ModelInfo): @classmethod @@ -85,7 +86,12 @@ DeepseekR1DistillQwenMeta = ModelMeta( model_sizes=["1.5", "7", "14", "32"], model_info_cls=DeepseekR1ModelInfo, is_multimodal=False, - quant_types=[QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF] + quant_types={ + "1.5": [QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF], + "7": [QuantType.UNSLOTH, QuantType.BNB], + "14": [QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF], + "32": [QuantType.GGUF, QuantType.BNB], + }, ) # "Qwen-7B-unsloth-bnb-4bit", @@ -98,45 +104,54 @@ DeepseekR1DistillQwenMeta = ModelMeta( # "Qwen-14B-unsloth-bnb-4bit", def register_deepseek_v3_models(include_original_model: bool = False): - global _IS_DEEPSEEKV3_REGISTERED - if _IS_DEEPSEEKV3_REGISTERED: + global _IS_DEEPSEEK_V3_REGISTERED + if _IS_DEEPSEEK_V3_REGISTERED: return _register_models(DeepseekV3Meta, include_original_model=include_original_model) - _register_models(DeepseekV3_0324Meta, include_original_model=include_original_model) - _IS_DEEPSEEKV3_REGISTERED = True + _IS_DEEPSEEK_V3_REGISTERED = True +def register_deepseek_v3_0324_models(include_original_model: bool = False): + global _IS_DEEPSEEK_V3_0324_REGISTERED + if _IS_DEEPSEEK_V3_0324_REGISTERED: + return + _register_models(DeepseekV3_0324Meta, include_original_model=include_original_model) + _IS_DEEPSEEK_V3_0324_REGISTERED = True def register_deepseek_r1_models(include_original_model: bool = False): - global _IS_DEEPSEEKR1_REGISTERED - if _IS_DEEPSEEKR1_REGISTERED: + global _IS_DEEPSEEK_R1_REGISTERED + if _IS_DEEPSEEK_R1_REGISTERED: return _register_models(DeepseekR1Meta, include_original_model=include_original_model) - _IS_DEEPSEEKR1_REGISTERED = True + _IS_DEEPSEEK_R1_REGISTERED = True def register_deepseek_r1_zero_models(include_original_model: bool = False): - global _IS_DEEPSEEKR1_ZERO_REGISTERED - if _IS_DEEPSEEKR1_ZERO_REGISTERED: + global _IS_DEEPSEEK_R1_ZERO_REGISTERED + if _IS_DEEPSEEK_R1_ZERO_REGISTERED: return _register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model) - _IS_DEEPSEEKR1_ZERO_REGISTERED = True + _IS_DEEPSEEK_R1_ZERO_REGISTERED = True def register_deepseek_r1_distill_llama_models(include_original_model: bool = False): - global _IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED - if _IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED: + global _IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED + if _IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED: return _register_models(DeepseekR1DistillLlamaMeta, include_original_model=include_original_model) - _IS_DEEPSEEKR1_DISTILL_LLAMA_REGISTERED = True + _IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED = True def register_deepseek_r1_distill_qwen_models(include_original_model: bool = False): - global _IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED - if _IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED: + global _IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED + if _IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED: return _register_models(DeepseekR1DistillQwenMeta, include_original_model=include_original_model) - _IS_DEEPSEEKR1_DISTILL_QWEN_REGISTERED = True + _IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED = True -def register_deepseek_r1_distill_models(include_original_model: bool = False): - register_deepseek_r1_distill_qwen_models(include_original_model=include_original_model) +def register_deepseek_models(include_original_model: bool = False): + register_deepseek_v3_models(include_original_model=include_original_model) + register_deepseek_v3_0324_models(include_original_model=include_original_model) + register_deepseek_r1_models(include_original_model=include_original_model) + register_deepseek_r1_zero_models(include_original_model=include_original_model) register_deepseek_r1_distill_llama_models(include_original_model=include_original_model) + register_deepseek_r1_distill_qwen_models(include_original_model=include_original_model) def _list_deepseek_r1_distill_models(): from unsloth.utils.hf_hub import ModelInfo as HfModelInfo @@ -153,9 +168,7 @@ def _list_deepseek_r1_distill_models(): return distill_models -register_deepseek_v3_models(include_original_model=True) -register_deepseek_r1_models(include_original_model=True) -register_deepseek_r1_distill_models(include_original_model=True) +register_deepseek_models(include_original_model=True) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info From 65ea6356e4b1529a9e4110bfc1539d57cc392880 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:08:11 -0700 Subject: [PATCH 29/38] refactor naming for mistral and phi --- unsloth/registry/_deepseek.py | 9 -------- unsloth/registry/_mistral.py | 15 +++++++------ unsloth/registry/_phi.py | 40 ++++++++++++++++++----------------- 3 files changed, 29 insertions(+), 35 deletions(-) diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 1f97a02f1c..854a62c00b 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -93,15 +93,6 @@ DeepseekR1DistillQwenMeta = ModelMeta( "32": [QuantType.GGUF, QuantType.BNB], }, ) - - # "Qwen-7B-unsloth-bnb-4bit", - # "Qwen-1.5B-unsloth-bnb-4bit", - # "Qwen-32B-GGUF", - - # "Qwen-14B-GGUF", - # "Qwen-32B-bnb-4bit", - # "Qwen-1.5B-GGUF", - # "Qwen-14B-unsloth-bnb-4bit", def register_deepseek_v3_models(include_original_model: bool = False): global _IS_DEEPSEEK_V3_REGISTERED diff --git a/unsloth/registry/_mistral.py b/unsloth/registry/_mistral.py index 65f1256708..c41b1f55b6 100644 --- a/unsloth/registry/_mistral.py +++ b/unsloth/registry/_mistral.py @@ -42,21 +42,22 @@ MistralSmall_2501_Base_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION MistralSmall_2501_Instruct_Meta = copy.deepcopy(MistralSmall_2503_Instruct_Meta) MistralSmall_2501_Instruct_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION -def register_mistral_small_models(): +def register_mistral_small_models(include_original_model: bool = False): global _IS_MISTRAL_SMALL_REGISTERED if _IS_MISTRAL_SMALL_REGISTERED: return - _register_models(MistralSmall_2503_Base_Meta) - _register_models(MistralSmall_2503_Instruct_Meta) - _register_models(MistralSmall_2501_Base_Meta) - _register_models(MistralSmall_2501_Instruct_Meta) + _register_models(MistralSmall_2503_Base_Meta, include_original_model=include_original_model) + _register_models(MistralSmall_2503_Instruct_Meta, include_original_model=include_original_model) + _register_models(MistralSmall_2501_Base_Meta, include_original_model=include_original_model) + _register_models(MistralSmall_2501_Instruct_Meta, include_original_model=include_original_model) _IS_MISTRAL_SMALL_REGISTERED = True -register_mistral_small_models() - +def register_mistral_models(include_original_model: bool = False): + register_mistral_small_models(include_original_model=include_original_model) if __name__ == "__main__": + register_mistral_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) diff --git a/unsloth/registry/_phi.py b/unsloth/registry/_phi.py index c69eaf83bb..9f23c494d5 100644 --- a/unsloth/registry/_phi.py +++ b/unsloth/registry/_phi.py @@ -1,7 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models -_IS_PHI_REGISTERED = False -_IS_PHI_INSTRUCT_REGISTERED = False +_IS_PHI_4_REGISTERED = False +_IS_PHI_4_INSTRUCT_REGISTERED = False class PhiModelInfo(ModelInfo): @classmethod @@ -10,7 +10,7 @@ class PhiModelInfo(ModelInfo): return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) # Phi Model Meta -PhiMeta = ModelMeta( +PhiMeta4 = ModelMeta( org="microsoft", base_name="phi", instruct_tags=[None], @@ -22,7 +22,7 @@ PhiMeta = ModelMeta( ) # Phi Instruct Model Meta -PhiInstructMeta = ModelMeta( +PhiInstructMeta4 = ModelMeta( org="microsoft", base_name="phi", instruct_tags=["mini-instruct"], @@ -33,24 +33,26 @@ PhiInstructMeta = ModelMeta( quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF], ) +def register_phi_4_models(include_original_model: bool = False): + global _IS_PHI_4_REGISTERED + if _IS_PHI_4_REGISTERED: + return + _register_models(PhiMeta4, include_original_model=include_original_model) + _IS_PHI_4_REGISTERED = True + +def register_phi_4_instruct_models(include_original_model: bool = False): + global _IS_PHI_4_INSTRUCT_REGISTERED + if _IS_PHI_4_INSTRUCT_REGISTERED: + return + _register_models(PhiInstructMeta4, include_original_model=include_original_model) + _IS_PHI_4_INSTRUCT_REGISTERED = True + def register_phi_models(include_original_model: bool = False): - global _IS_PHI_REGISTERED - if _IS_PHI_REGISTERED: - return - _register_models(PhiMeta, include_original_model=include_original_model) - _IS_PHI_REGISTERED = True - -def register_phi_instruct_models(include_original_model: bool = False): - global _IS_PHI_INSTRUCT_REGISTERED - if _IS_PHI_INSTRUCT_REGISTERED: - return - _register_models(PhiInstructMeta, include_original_model=include_original_model) - _IS_PHI_INSTRUCT_REGISTERED = True - -register_phi_models(include_original_model=True) -register_phi_instruct_models(include_original_model=True) + register_phi_4_models(include_original_model=include_original_model) + register_phi_4_instruct_models(include_original_model=include_original_model) if __name__ == "__main__": + register_phi_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) From 2ff490e23b2015a42a2e50be3f1f6fdc485b8332 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:11:35 -0700 Subject: [PATCH 30/38] add global register models --- unsloth/registry/__init__.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py index dd5b45c4ee..154cea6deb 100644 --- a/unsloth/registry/__init__.py +++ b/unsloth/registry/__init__.py @@ -1,5 +1,13 @@ -# from ._deepseek import register_deepseek_models, register -# from ._llama import register_llama_models, register_llama_vision_models -# from ._mistral import register_mistral_models -# from ._openai import register_openai_models -# from ._qwen import register_qwen_models +from ._deepseek import register_deepseek_models as _register_deepseek_models +from ._gemma import register_gemma_models as _register_gemma_models +from ._llama import register_llama_models as _register_llama_models +from ._mistral import register_mistral_models as _register_mistral_models +from ._phi import register_phi_models as _register_phi_models +from ._qwen import register_qwen_models as _register_qwen_models + +_register_deepseek_models() +_register_gemma_models() +_register_llama_models() +_register_mistral_models() +_register_phi_models() +_register_qwen_models() \ No newline at end of file From d93120db9dee07a413911d8f72fbb1086477ab5f Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:22:26 -0700 Subject: [PATCH 31/38] refactor model registration tests for new registry apis --- tests/test_model_registry.py | 89 +++++++++++++++++++----------------- unsloth/registry/__init__.py | 15 +++--- 2 files changed, 55 insertions(+), 49 deletions(-) diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 183edc92d5..1f9ddd922e 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -2,39 +2,39 @@ from dataclasses import dataclass import pytest from huggingface_hub import ModelInfo as HfModelInfo -from unsloth.model_registry import ( - ModelInfo, - get_llama_models, - get_llama_vision_models, - get_phi_instruct_models, - get_phi_models, - get_qwen_models, - get_qwen_vl_models, -) + +from unsloth.registry import register_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, ModelInfo from unsloth.utils.hf_hub import get_model_info MODEL_NAMES = [ "llama", - "llama_vision", "qwen", - "qwen_vl", + "mistral", "phi", - "phi_instruct", + "gemma", + "deepseek", ] -REGISTERED_MODELS = [ - get_llama_models(), - get_llama_vision_models(), - get_qwen_models(), - get_qwen_vl_models(), - get_phi_models(), - get_phi_instruct_models(), +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 - models: dict[str, ModelInfo] + registration_models: callable def _test_model_uploaded(model_ids: list[str]): @@ -49,37 +49,40 @@ def _test_model_uploaded(model_ids: list[str]): TestParams = [ ModelTestParam(name, models) - for name, models in zip(MODEL_NAMES, REGISTERED_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_uploaded(model_test_param: ModelTestParam): - missing_models = _test_model_uploaded(model_test_param.models) +def test_model_registration(model_test_param: ModelTestParam): + MODEL_REGISTRY.clear() + model_test_param.registration_models() + 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}" ) -if __name__ == "__main__": - for method in [ - get_llama_models, - get_llama_vision_models, - get_qwen_models, - get_qwen_vl_models, - get_phi_models, - get_phi_instruct_models, - ]: - models = method() - model_name = next(iter(models.values())).base_name - print(f"{model_name}: {len(models)} registered") - for model_info in models.values(): - print(f" {model_info.model_path}") - missing_models = test_model_uploaded(list(models.keys())) +# if __name__ == "__main__": +# for method in [ +# get_llama_models, +# get_llama_vision_models, +# get_qwen_models, +# get_qwen_vl_models, +# get_phi_models, +# get_phi_instruct_models, +# ]: +# models = method() +# model_name = next(iter(models.values())).base_name +# print(f"{model_name}: {len(models)} registered") +# for model_info in models.values(): +# print(f" {model_info.model_path}") +# missing_models = test_model_uploaded(list(models.keys())) - if missing_models: - print("--------------------------------") - print(f"Missing models: {missing_models}") - print("--------------------------------") +# if missing_models: +# print("--------------------------------") +# print(f"Missing models: {missing_models}") +# print("--------------------------------") diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py index 154cea6deb..1b92fef74d 100644 --- a/unsloth/registry/__init__.py +++ b/unsloth/registry/__init__.py @@ -5,9 +5,12 @@ from ._mistral import register_mistral_models as _register_mistral_models from ._phi import register_phi_models as _register_phi_models from ._qwen import register_qwen_models as _register_qwen_models -_register_deepseek_models() -_register_gemma_models() -_register_llama_models() -_register_mistral_models() -_register_phi_models() -_register_qwen_models() \ No newline at end of file + +def register_models(): + _register_deepseek_models() + _register_gemma_models() + _register_llama_models() + _register_mistral_models() + _register_phi_models() + _register_qwen_models() + From 959727a2d214719394b14e7b570571d96ea18565 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:36:19 -0700 Subject: [PATCH 32/38] add model search method --- tests/test_model_registry.py | 47 +++++++++++++----------------- unsloth/registry/__init__.py | 37 ++++++++++++++++++++++- unsloth/registry/model_registry.py | 2 +- 3 files changed, 58 insertions(+), 28 deletions(-) diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 1f9ddd922e..a767d42cdb 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -1,3 +1,13 @@ +""" + +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 @@ -10,7 +20,7 @@ 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, ModelInfo +from unsloth.registry.registry import MODEL_REGISTRY from unsloth.utils.hf_hub import get_model_info MODEL_NAMES = [ @@ -34,7 +44,7 @@ MODEL_REGISTRATION_METHODS = [ @dataclass class ModelTestParam: name: str - registration_models: callable + register_models: callable def _test_model_uploaded(model_ids: list[str]): @@ -52,13 +62,13 @@ TestParams = [ 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 -) +@pytest.mark.parametrize("model_test_param", TestParams, ids=lambda param: param.name) def test_model_registration(model_test_param: ModelTestParam): MODEL_REGISTRY.clear() - model_test_param.registration_models() + 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, ( @@ -66,23 +76,8 @@ def test_model_registration(model_test_param: ModelTestParam): ) -# if __name__ == "__main__": -# for method in [ -# get_llama_models, -# get_llama_vision_models, -# get_qwen_models, -# get_qwen_vl_models, -# get_phi_models, -# get_phi_instruct_models, -# ]: -# models = method() -# model_name = next(iter(models.values())).base_name -# print(f"{model_name}: {len(models)} registered") -# for model_info in models.values(): -# print(f" {model_info.model_path}") -# missing_models = test_model_uploaded(list(models.keys())) - -# if missing_models: -# print("--------------------------------") -# print(f"Missing models: {missing_models}") -# print("--------------------------------") +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}" diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py index 1b92fef74d..a46ab773d8 100644 --- a/unsloth/registry/__init__.py +++ b/unsloth/registry/__init__.py @@ -4,9 +4,15 @@ from ._llama import register_llama_models as _register_llama_models from ._mistral import register_mistral_models as _register_mistral_models from ._phi import register_phi_models as _register_phi_models from ._qwen import register_qwen_models as _register_qwen_models +from .registry import MODEL_REGISTRY, ModelInfo, QuantType +_ARE_MODELS_REGISTERED = False -def register_models(): +def register_models(): + global _ARE_MODELS_REGISTERED + + if _ARE_MODELS_REGISTERED: + return _register_deepseek_models() _register_gemma_models() _register_llama_models() @@ -14,3 +20,32 @@ def register_models(): _register_phi_models() _register_qwen_models() + _ARE_MODELS_REGISTERED = True + +def get_model_info(org: str = None, base_name: str = None, version: str = None, size: str = None, quant_types: list[QuantType] = None, search_pattern: str = None) -> list[ModelInfo]: + """ + Get model info from the registry. + + See registry.ModelInfo for more fields. + + If search_pattern is provided, the full model path will be matched against the pattern, where the model path is the model_id on huggingface hub. + + """ + if not _ARE_MODELS_REGISTERED: + register_models() + + model_infos = MODEL_REGISTRY.values() + if org: + model_infos = [model_info for model_info in model_infos if model_info.org == org] + if base_name: + model_infos = [model_info for model_info in model_infos if model_info.base_name == base_name] + if version: + model_infos = [model_info for model_info in model_infos if model_info.version == version] + if size: + model_infos = [model_info for model_info in model_infos if model_info.size == size] + if quant_types: + model_infos = [model_info for model_info in model_infos if any(model_info.quant_type == quant_type for quant_type in quant_types)] + if search_pattern: + model_infos = [model_info for model_info in model_infos if search_pattern in model_info.model_path] + + return model_infos \ No newline at end of file diff --git a/unsloth/registry/model_registry.py b/unsloth/registry/model_registry.py index de9609934c..b51644beb7 100644 --- a/unsloth/registry/model_registry.py +++ b/unsloth/registry/model_registry.py @@ -306,4 +306,4 @@ if __name__ == "__main__": if len(missing_models) == 0: # print unicode checkmark - print(f"\u2713 All models found!") \ No newline at end of file + print("\u2713 All models found!") \ No newline at end of file From bb66d454e2a0c91ac2b635fe74fd80d556552e1e Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:36:42 -0700 Subject: [PATCH 33/38] remove deprecated registration api --- unsloth/registry/model_registry.py | 309 ----------------------------- 1 file changed, 309 deletions(-) delete mode 100644 unsloth/registry/model_registry.py diff --git a/unsloth/registry/model_registry.py b/unsloth/registry/model_registry.py deleted file mode 100644 index b51644beb7..0000000000 --- a/unsloth/registry/model_registry.py +++ /dev/null @@ -1,309 +0,0 @@ -from functools import partial -from typing import Callable, Literal - -from unsloth.registry._llama import LlamaMeta3_1, LlamaMeta3_2 -from unsloth.registry.common import ModelInfo, ModelMeta - -# _IS_LLAMA_REGISTERED = False -# _IS_LLAMA_VISION_REGISTERED = False - -# _IS_QWEN_REGISTERED = False -# _IS_QWEN_VL_REGISTERED = False - -_IS_GEMMA_REGISTERED = False - -_IS_PHI_REGISTERED = False -_IS_PHI_INSTRUCT_REGISTERED = False - - - -# class PhiModelInfo(ModelInfo): -# @classmethod -# def construct_model_name( -# cls, base_name, version, size, quant_type, instruct_tag -# ): -# key = f"{base_name}-{version}" -# key = cls.append_instruct_tag(key, instruct_tag) -# key = cls.append_quant_type(key, quant_type) -# return key - - - - - -# # Qwen text only models -# # NOTE: Qwen vision models will be registered separately - -# _PHI_INFO = { -# "org": "microsoft", -# "base_name": "phi", -# "model_versions": ["4"], -# "model_sizes": {"4": [None]}, # -1 means only 1 size -# "instruct_tags": [None], -# "is_multimodal": False, -# "model_info_cls": PhiModelInfo, -# } - -# _PHI_INSTRUCT_INFO = { -# "org": "microsoft", -# "base_name": "Phi", -# "model_versions": ["4"], -# "model_sizes": {"4": [None]}, # -1 means only 1 size -# "instruct_tags": ["mini-instruct"], -# "is_multimodal": False, -# "model_info_cls": PhiModelInfo, -# } - - -MODEL_REGISTRY: dict[str, ModelInfo] = {} - - -def register_model( - model_info_cls: ModelInfo, - org: str, - base_name: str, - version: str, - size: int, - instruct_tag: str = None, - quant_type: Literal["bnb", "unsloth"] = None, - is_multimodal: bool = False, - name: str = None, -): - name = name or model_info_cls.construct_model_name( - base_name=base_name, - version=version, - size=size, - quant_type=quant_type, - instruct_tag=instruct_tag, - ) - key = f"{org}/{name}" - - if key in MODEL_REGISTRY: - raise ValueError(f"Model {key} already registered") - - MODEL_REGISTRY[key] = model_info_cls( - org=org, - base_name=base_name, - version=version, - size=size, - is_multimodal=is_multimodal, - instruct_tag=instruct_tag, - quant_type=quant_type, - name=name, - ) - - -# def _register_models(model_info: dict): -# org = model_info["org"] -# base_name = model_info["base_name"] -# instruct_tags = model_info["instruct_tags"] -# model_versions = model_info["model_versions"] -# model_sizes = model_info["model_sizes"] -# is_multimodal = model_info["is_multimodal"] -# model_info_cls = model_info["model_info_cls"] - -# for version in model_versions: -# for size in model_sizes[version]: -# for instruct_tag in instruct_tags: -# for quant_type in QUANT_TYPES: -# _org = "unsloth" if quant_type is not None else org -# register_model( -# model_info_cls=model_info_cls, -# org=_org, -# base_name=base_name, -# version=version, -# size=size, -# instruct_tag=instruct_tag, -# quant_type=quant_type, -# is_multimodal=is_multimodal, -# ) - - -def _register_models(model_meta: ModelMeta): - org = model_meta.org - base_name = model_meta.base_name - instruct_tags = model_meta.instruct_tags - model_version = model_meta.model_version - model_sizes = model_meta.model_sizes - is_multimodal = model_meta.is_multimodal - quant_types = model_meta.quant_types - model_info_cls = model_meta.model_info_cls - - for size in model_sizes: - for instruct_tag in instruct_tags: - for quant_type in quant_types: - _org = "unsloth" if quant_type is not None else org - register_model( - model_info_cls=model_info_cls, - org=_org, - base_name=base_name, - version=model_version, - size=size, - instruct_tag=instruct_tag, - quant_type=quant_type, - is_multimodal=is_multimodal, - ) - -def register_llama_models(): - global _IS_LLAMA_REGISTERED - if _IS_LLAMA_REGISTERED: - return - _register_models(LlamaMeta3_1) - _register_models(LlamaMeta3_2) - _IS_LLAMA_REGISTERED = True - - -def register_llama_vision_models(): - global _IS_LLAMA_VISION_REGISTERED - if _IS_LLAMA_VISION_REGISTERED: - return - _register_models(_LLAMA_VISION_INFO) - _IS_LLAMA_VISION_REGISTERED = True - - -def register_qwen_models(): - global _IS_QWEN_REGISTERED - if _IS_QWEN_REGISTERED: - return - - _register_models(_QWEN_INFO) - _IS_QWEN_REGISTERED = True - - -def register_qwen_vl_models(): - global _IS_QWEN_VL_REGISTERED - if _IS_QWEN_VL_REGISTERED: - return - - _register_models(_QWEN_VL_INFO) - _IS_QWEN_VL_REGISTERED = True - - -def register_gemma_models(): - global _IS_GEMMA_REGISTERED - _register_models(_GEMMA_INFO) - _IS_GEMMA_REGISTERED = True - - -def register_phi_models(): - global _IS_PHI_REGISTERED - if _IS_PHI_REGISTERED: - return - _register_models(_PHI_INFO) - _IS_PHI_REGISTERED = True - - -def register_phi_instruct_models(): - global _IS_PHI_INSTRUCT_REGISTERED - if _IS_PHI_INSTRUCT_REGISTERED: - return - - _register_models(_PHI_INSTRUCT_INFO) - _IS_PHI_INSTRUCT_REGISTERED = True - - -def _base_name_filter(model_info: ModelInfo, base_name: str): - return model_info.base_name == base_name - - -def _get_models(filter_func: Callable[[ModelInfo], bool] = _base_name_filter): - return {k: v for k, v in MODEL_REGISTRY.items() if filter_func(v)} - - -def get_llama_models(version: str = None): - if not _IS_LLAMA_REGISTERED: - register_llama_models() - - llama_models: dict[str, ModelInfo] = _get_models( - partial(_base_name_filter, base_name=LlamaMeta3_1.base_name) - ) - if version is not None: - llama_models = { - k: v for k, v in llama_models.items() if v.version == version - } - return llama_models - - -def get_llama_vision_models(): - if not _IS_LLAMA_VISION_REGISTERED: - register_llama_vision_models() - - return _get_models( - lambda model_info: model_info.base_name - == _LLAMA_VISION_INFO["base_name"] - and model_info.is_multimodal - ) - - -def get_qwen_models(): - if not _IS_QWEN_REGISTERED: - register_qwen_models() - - return _get_models( - lambda model_info: model_info.base_name == _QWEN_INFO["base_name"] - ) - - -def get_qwen_vl_models(): - if not _IS_QWEN_VL_REGISTERED: - register_qwen_vl_models() - return _get_models( - lambda model_info: model_info.base_name == _QWEN_VL_INFO["base_name"] - ) - - -def get_gemma_models(): - if not _IS_GEMMA_REGISTERED: - register_gemma_models() - - return _get_models( - lambda model_info: model_info.base_name == _GEMMA_INFO["base_name"] - ) - - -def get_phi_models(): - if not _IS_PHI_REGISTERED: - register_phi_models() - return _get_models( - lambda model_info: model_info.base_name == _PHI_INFO["base_name"] - ) - - -def get_phi_instruct_models(): - if not _IS_PHI_INSTRUCT_REGISTERED: - register_phi_instruct_models() - return _get_models( - lambda model_info: model_info.base_name - == _PHI_INSTRUCT_INFO["base_name"] - ) - - -if __name__ == "__main__": - from huggingface_hub import HfApi - - api = HfApi() - - def get_model_info( - model_id: str, properties: list[str] = None - ) -> ModelInfo: - try: - model_info: ModelInfo = api.model_info(model_id, expand=properties) - except Exception as e: - print(f"Error getting model info for {model_id}: {e}") - model_info = None - return model_info - - register_llama_models() - - llama3_1_models = get_llama_models(version="3.2") - missing_models = [] - for k, v in llama3_1_models.items(): - model_info = get_model_info(v.model_path) - if model_info is None: - # print unicode cross mark followed by model k - print(f"\u2718 {k}") - missing_models.append(k) - - if len(missing_models) == 0: - # print unicode checkmark - print("\u2713 All models found!") \ No newline at end of file From ecf70d6caa1b2b3d2b55d3773269e8c5e0e5050d Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 17:58:44 -0700 Subject: [PATCH 34/38] add quant type test --- tests/test_model_registry.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index a767d42cdb..3d570af230 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -13,15 +13,14 @@ from dataclasses import dataclass import pytest from huggingface_hub import ModelInfo as HfModelInfo -from unsloth.registry import register_models +from unsloth.registry import get_model_info, register_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 -from unsloth.utils.hf_hub import get_model_info +from unsloth.registry.registry import MODEL_REGISTRY, QUANT_TAG_MAP, QuantType MODEL_NAMES = [ "llama", @@ -81,3 +80,11 @@ def test_all_model_registration(): 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 = get_model_info(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) \ No newline at end of file From e2cfec6339a502a2af4bf5e03cd581daf7fdb735 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 18:11:58 -0700 Subject: [PATCH 35/38] add registry readme --- unsloth/registry/REGISTRY.md | 45 ++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 unsloth/registry/REGISTRY.md diff --git a/unsloth/registry/REGISTRY.md b/unsloth/registry/REGISTRY.md new file mode 100644 index 0000000000..b794d26be6 --- /dev/null +++ b/unsloth/registry/REGISTRY.md @@ -0,0 +1,45 @@ +## Model Registry + +### Structure + +Each model is registered in a separate file within the `registry` module (e.g. `registry/_llama.py`). + +Within each model registration file, a high-level `ModelMeta` is created for each model version, with the following structure: +```python +@dataclass +class ModelMeta: + org: str + base_name: str + model_version: str + model_info_cls: type[ModelInfo] + model_sizes: list[str] = field(default_factory=list) + instruct_tags: list[str] = field(default_factory=list) + quant_types: list[QuantType] | dict[str, list[QuantType]] = field(default_factory=list) + is_multimodal: bool = False +``` + +Each model then instantiates a global `ModelMeta` for its specific model version, defining how the model path (e.g. `unsloth/Llama-3.1-8B-Instruct`) is constructed since each model type has a different naming convention. +```python +LlamaMeta_3_1 = ModelMeta( + org="meta-llama", + base_name="Llama", + instruct_tags=[None, "Instruct"], + model_version="3.1", + model_sizes=["8"], + model_info_cls=LlamaModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH], +) +``` + +`LlamaModelInfo` is a subclass of `ModelInfo` that defines the model path for each model size and quant type. +```python +class LlamaModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}-{size}B" + return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) +``` + +Once these constructs are defined, the model is registered in the `registry` module by calling `register_models` with the `ModelMeta` and `ModelInfo` classes. + From 8d393f29c10d5700a37338774499cc38e3215cd5 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 18:12:52 -0700 Subject: [PATCH 36/38] make llama registration more specific --- unsloth/registry/_llama.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index c84c5b8d30..ec6e39a86d 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -1,6 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models -_IS_LLAMA_3_REGISTERED = False +_IS_LLAMA_3_1_REGISTERED = False +_IS_LLAMA_3_2_REGISTERED = False _IS_LLAMA_3_2_VISION_REGISTERED = False @@ -70,14 +71,20 @@ LlamaMeta_3_2_Vision = ModelMeta( ) -def register_llama_3_models(include_original_model: bool = False): - global _IS_LLAMA_3_REGISTERED - if _IS_LLAMA_3_REGISTERED: +def register_llama_3_1_models(include_original_model: bool = False): + global _IS_LLAMA_3_1_REGISTERED + if _IS_LLAMA_3_1_REGISTERED: return _register_models(LlamaMeta_3_1, include_original_model=include_original_model) + _IS_LLAMA_3_1_REGISTERED = True + +def register_llama_3_2_models(include_original_model: bool = False): + global _IS_LLAMA_3_2_REGISTERED + if _IS_LLAMA_3_2_REGISTERED: + return _register_models(LlamaMeta_3_2_Base, include_original_model=include_original_model) _register_models(LlamaMeta_3_2_Instruct, include_original_model=include_original_model) - _IS_LLAMA_3_REGISTERED = True + _IS_LLAMA_3_2_REGISTERED = True def register_llama_3_2_vision_models(include_original_model: bool = False): global _IS_LLAMA_3_2_VISION_REGISTERED @@ -88,7 +95,8 @@ def register_llama_3_2_vision_models(include_original_model: bool = False): def register_llama_models(include_original_model: bool = False): - register_llama_3_models(include_original_model=include_original_model) + register_llama_3_1_models(include_original_model=include_original_model) + register_llama_3_2_models(include_original_model=include_original_model) register_llama_3_2_vision_models(include_original_model=include_original_model) if __name__ == "__main__": From b33970525c44ff6cef3f497c3fde0548c5b0d449 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 18:24:15 -0700 Subject: [PATCH 37/38] clear registry when executing individual model registration file --- tests/test_model_registry.py | 5 ++-- unsloth/registry/REGISTRY.md | 50 ++++++++++++++++++++++++++++++++++- unsloth/registry/__init__.py | 2 +- unsloth/registry/_deepseek.py | 4 +++ unsloth/registry/_gemma.py | 5 +++- unsloth/registry/_llama.py | 4 ++- unsloth/registry/_mistral.py | 5 +++- unsloth/registry/_phi.py | 5 +++- unsloth/registry/_qwen.py | 5 +++- 9 files changed, 76 insertions(+), 9 deletions(-) diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 3d570af230..f59f4f0dab 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -13,7 +13,7 @@ from dataclasses import dataclass import pytest from huggingface_hub import ModelInfo as HfModelInfo -from unsloth.registry import get_model_info, register_models +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 @@ -21,6 +21,7 @@ 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", @@ -84,7 +85,7 @@ def test_all_model_registration(): 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 = get_model_info(quant_types=[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) \ No newline at end of file diff --git a/unsloth/registry/REGISTRY.md b/unsloth/registry/REGISTRY.md index b794d26be6..8240d686e6 100644 --- a/unsloth/registry/REGISTRY.md +++ b/unsloth/registry/REGISTRY.md @@ -1,6 +1,16 @@ ## Model Registry ### Structure +``` +unsloth + -registry + __init__.py + registry.py + _llama.py + _mistral.py + _phi.py + ... +``` Each model is registered in a separate file within the `registry` module (e.g. `registry/_llama.py`). @@ -41,5 +51,43 @@ class LlamaModelInfo(ModelInfo): return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) ``` -Once these constructs are defined, the model is registered in the `registry` module by calling `register_models` with the `ModelMeta` and `ModelInfo` classes. +Once these constructs are defined, the model is registered by writing a register_xx_models function. +```python +def register_llama_3_1_models(include_original_model: bool = False): + global _IS_LLAMA_3_1_REGISTERED + if _IS_LLAMA_3_1_REGISTERED: + return + _register_models(LlamaMeta_3_1, include_original_model=include_original_model) + _IS_LLAMA_3_1_REGISTERED = True +``` + +`_register_models` is a helper function that registers the model with the registry. The global `_IS_XX_REGISTERED` is used to prevent duplicate registration. + +Once a model is registered, registry.registry.MODEL_REGISTRY is updated with the model info and can be searched with `registry.search_models`. + +### Tests + +The `tests/test_model_registry.py` file contains tests for the model registry. + +Also, each model registration file is an executable module that checks that all registered models are available on `huggingface_hub`. +```python +python unsloth.registry._llama.py +``` + +Prints the following (abridged) output: +```bash +✓ unsloth/Llama-3.1-8B +✓ unsloth/Llama-3.1-8B-bnb-4bit +✓ unsloth/Llama-3.1-8B-unsloth-bnb-4bit +✓ meta-llama/Llama-3.1-8B +✓ unsloth/Llama-3.1-8B-Instruct +✓ unsloth/Llama-3.1-8B-Instruct-bnb-4bit +✓ unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit +✓ meta-llama/Llama-3.1-8B-Instruct +✓ unsloth/Llama-3.2-1B +✓ unsloth/Llama-3.2-1B-bnb-4bit +✓ unsloth/Llama-3.2-1B-unsloth-bnb-4bit +✓ meta-llama/Llama-3.2-1B +... +``` diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py index a46ab773d8..5874743694 100644 --- a/unsloth/registry/__init__.py +++ b/unsloth/registry/__init__.py @@ -22,7 +22,7 @@ def register_models(): _ARE_MODELS_REGISTERED = True -def get_model_info(org: str = None, base_name: str = None, version: str = None, size: str = None, quant_types: list[QuantType] = None, search_pattern: str = None) -> list[ModelInfo]: +def search_models(org: str = None, base_name: str = None, version: str = None, size: str = None, quant_types: list[QuantType] = None, search_pattern: str = None) -> list[ModelInfo]: """ Get model info from the registry. diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 854a62c00b..153a0e508e 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -163,6 +163,10 @@ register_deepseek_models(include_original_model=True) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_deepseek_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_gemma.py b/unsloth/registry/_gemma.py index 8c47e7e69d..9490c84f2f 100644 --- a/unsloth/registry/_gemma.py +++ b/unsloth/registry/_gemma.py @@ -53,8 +53,11 @@ def register_gemma_models(include_original_model: bool = False): if __name__ == "__main__": - register_gemma_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_gemma_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index ec6e39a86d..1c2dd5bf18 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -100,8 +100,10 @@ def register_llama_models(include_original_model: bool = False): register_llama_3_2_vision_models(include_original_model=include_original_model) if __name__ == "__main__": - register_llama_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_llama_models(include_original_model=True) for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) diff --git a/unsloth/registry/_mistral.py b/unsloth/registry/_mistral.py index c41b1f55b6..44cd1e7646 100644 --- a/unsloth/registry/_mistral.py +++ b/unsloth/registry/_mistral.py @@ -57,8 +57,11 @@ def register_mistral_models(include_original_model: bool = False): register_mistral_small_models(include_original_model=include_original_model) if __name__ == "__main__": - register_mistral_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_mistral_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_phi.py b/unsloth/registry/_phi.py index 9f23c494d5..d06ec8d377 100644 --- a/unsloth/registry/_phi.py +++ b/unsloth/registry/_phi.py @@ -52,8 +52,11 @@ def register_phi_models(include_original_model: bool = False): register_phi_4_instruct_models(include_original_model=include_original_model) if __name__ == "__main__": - register_phi_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_phi_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py index c364f9b099..4417515a77 100644 --- a/unsloth/registry/_qwen.py +++ b/unsloth/registry/_qwen.py @@ -104,8 +104,11 @@ def register_qwen_models(include_original_model: bool = False): register_qwen_qwq_models(include_original_model=include_original_model) if __name__ == "__main__": - register_qwen_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_qwen_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: From 9a14edcd2f760bdb8beaa945549d87afb0b9684c Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 18:34:18 -0700 Subject: [PATCH 38/38] more registry readme updates --- unsloth/registry/REGISTRY.md | 17 +++++++++++++++++ unsloth/registry/_llama.py | 2 +- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/unsloth/registry/REGISTRY.md b/unsloth/registry/REGISTRY.md index 8240d686e6..a0b3d96cad 100644 --- a/unsloth/registry/REGISTRY.md +++ b/unsloth/registry/REGISTRY.md @@ -91,3 +91,20 @@ Prints the following (abridged) output: ... ``` +### TODO +- Model Collections + - [x] Gemma3 + - [ ] Llama3.1 + - [x] Llama3.2 + - [x] MistralSmall + - [x] Qwen2.5 + - [x] Qwen2.5-VL + - [ ] Qwen2.5 Coder + - [x] QwenQwQ-32B + - [x] Deepseek v3 + - [x] Deepseek R1 + - [x] Phi-4 + - [ ] Unsloth 4-bit Dynamic Quants + - [ ] Vision/multimodal models +- Sync model uploads with registry +- Add utility methods for tracking model stats \ No newline at end of file diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index 1c2dd5bf18..f1b9dbdd32 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -102,7 +102,7 @@ def register_llama_models(include_original_model: bool = False): if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info MODEL_REGISTRY.clear() - + register_llama_models(include_original_model=True) for model_id, model_info in MODEL_REGISTRY.items():