From 6f6a7a5e9bb67ecd23e6486bf79160b485c35d82 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Fri, 28 Mar 2025 16:49:12 -0700 Subject: [PATCH] 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