From 73ac5bf9fbd6fdb3c835634accf02ae50ef3fbda Mon Sep 17 00:00:00 2001 From: jeromeku Date: Sun, 30 Mar 2025 11:36:48 -0700 Subject: [PATCH] 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, + )