Merge pull request #2255 from jeromeku/registry-refactor

Registry refactor
This commit is contained in:
Michael Han 2025-04-01 23:16:27 -07:00 committed by GitHub
commit fd20192aef
15 changed files with 1306 additions and 0 deletions

177
.gitignore vendored Normal file
View file

@ -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

View file

@ -32,6 +32,10 @@ include-package-data = false
exclude = ["images*", "tests*"]
[project.optional-dependencies]
dev = [
"pytest",
]
triton = [
"triton-windows ; platform_system == 'Windows'",
]

0
tests/__init__.py Normal file
View file

View file

@ -0,0 +1,91 @@
"""
Test model registration methods
Checks that model registration methods work for respective models as well as all models
The check is performed
- by registering the models
- checking that the instantiated models can be found on huggingface hub by querying for the model id
"""
from dataclasses import dataclass
import pytest
from huggingface_hub import ModelInfo as HfModelInfo
from unsloth.registry import register_models, search_models
from unsloth.registry._deepseek import register_deepseek_models
from unsloth.registry._gemma import register_gemma_models
from unsloth.registry._llama import register_llama_models
from unsloth.registry._mistral import register_mistral_models
from unsloth.registry._phi import register_phi_models
from unsloth.registry._qwen import register_qwen_models
from unsloth.registry.registry import MODEL_REGISTRY, QUANT_TAG_MAP, QuantType
from unsloth.utils.hf_hub import get_model_info
MODEL_NAMES = [
"llama",
"qwen",
"mistral",
"phi",
"gemma",
"deepseek",
]
MODEL_REGISTRATION_METHODS = [
register_llama_models,
register_qwen_models,
register_mistral_models,
register_phi_models,
register_gemma_models,
register_deepseek_models,
]
@dataclass
class ModelTestParam:
name: str
register_models: callable
def _test_model_uploaded(model_ids: list[str]):
missing_models = []
for _id in model_ids:
model_info: HfModelInfo = get_model_info(_id)
if not model_info:
missing_models.append(_id)
return missing_models
TestParams = [
ModelTestParam(name, models)
for name, models in zip(MODEL_NAMES, MODEL_REGISTRATION_METHODS)
]
# Test that model registration methods register respective models
@pytest.mark.parametrize("model_test_param", TestParams, ids=lambda param: param.name)
def test_model_registration(model_test_param: ModelTestParam):
MODEL_REGISTRY.clear()
registration_method = model_test_param.register_models
registration_method()
registered_models = MODEL_REGISTRY.keys()
missing_models = _test_model_uploaded(registered_models)
assert not missing_models, (
f"{model_test_param.name} missing following models: {missing_models}"
)
def test_all_model_registration():
register_models()
registered_models = MODEL_REGISTRY.keys()
missing_models = _test_model_uploaded(registered_models)
assert not missing_models, f"Missing following models: {missing_models}"
def test_quant_type():
# Test that the quant_type is correctly set for model paths
# NOTE: for models registered under org="unsloth" with QuantType.NONE aliases QuantType.UNSLOTH
dynamic_quant_models = search_models(quant_types=[QuantType.UNSLOTH])
assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models)
quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH]
assert all(quant_tag in m.model_path for m in dynamic_quant_models)

View file

@ -0,0 +1,110 @@
## 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`).
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 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
...
```
### 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

View file

@ -0,0 +1,51 @@
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
from .registry import MODEL_REGISTRY, ModelInfo, QuantType
_ARE_MODELS_REGISTERED = False
def register_models():
global _ARE_MODELS_REGISTERED
if _ARE_MODELS_REGISTERED:
return
_register_deepseek_models()
_register_gemma_models()
_register_llama_models()
_register_mistral_models()
_register_phi_models()
_register_qwen_models()
_ARE_MODELS_REGISTERED = True
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.
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

View file

@ -0,0 +1,179 @@
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
_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
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
key = f"{base_name}-V{version}"
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, 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
if size:
key = f"{key}-{size}B"
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, 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],
)
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],
)
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],
)
DeepseekR1DistillLlamaMeta = 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]},
)
# 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={
"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],
},
)
def register_deepseek_v3_models(include_original_model: bool = False):
global _IS_DEEPSEEK_V3_REGISTERED
if _IS_DEEPSEEK_V3_REGISTERED:
return
_register_models(DeepseekV3Meta, include_original_model=include_original_model)
_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_DEEPSEEK_R1_REGISTERED
if _IS_DEEPSEEK_R1_REGISTERED:
return
_register_models(DeepseekR1Meta, include_original_model=include_original_model)
_IS_DEEPSEEK_R1_REGISTERED = True
def register_deepseek_r1_zero_models(include_original_model: bool = False):
global _IS_DEEPSEEK_R1_ZERO_REGISTERED
if _IS_DEEPSEEK_R1_ZERO_REGISTERED:
return
_register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model)
_IS_DEEPSEEK_R1_ZERO_REGISTERED = True
def register_deepseek_r1_distill_llama_models(include_original_model: bool = False):
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_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED = True
def register_deepseek_r1_distill_qwen_models(include_original_model: bool = False):
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_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED = True
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
from unsloth.utils.hf_hub import list_models
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-")
distill_models.append(version)
return distill_models
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:
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)

View file

@ -0,0 +1,66 @@
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
_IS_GEMMA_3_BASE_REGISTERED = False
_IS_GEMMA_3_INSTRUCT_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"
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, 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_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)
_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)
if __name__ == "__main__":
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:
print(f"\u2718 {model_id}")
else:
print(f"\u2713 {model_id}")

113
unsloth/registry/_llama.py Normal file
View file

@ -0,0 +1,113 @@
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
_IS_LLAMA_3_1_REGISTERED = False
_IS_LLAMA_3_2_REGISTERED = False
_IS_LLAMA_3_2_VISION_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"
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"
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
# Llama 3.1
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],
)
# Llama 3.2 Base Models
LlamaMeta_3_2_Base = ModelMeta(
org="meta-llama",
base_name="Llama",
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
LlamaMeta_3_2_Instruct = ModelMeta(
org="meta-llama",
base_name="Llama",
instruct_tags=["Instruct"],
model_version="3.2",
model_sizes=["1", "3"],
model_info_cls=LlamaModelInfo,
is_multimodal=False,
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
)
# Llama 3.2 Vision
LlamaMeta_3_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={
"11": [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
"90": [QuantType.NONE],
},
)
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_2_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):
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 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)
if model_info is None:
print(f"\u2718 {model_id}")
else:
print(f"\u2713 {model_id}")

View file

@ -0,0 +1,70 @@
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(include_original_model: bool = False):
global _IS_MISTRAL_SMALL_REGISTERED
if _IS_MISTRAL_SMALL_REGISTERED:
return
_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
def register_mistral_models(include_original_model: bool = False):
register_mistral_small_models(include_original_model=include_original_model)
if __name__ == "__main__":
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:
print(f"\u2718 {model_id}")
else:
print(f"\u2713 {model_id}")

65
unsloth/registry/_phi.py Normal file
View file

@ -0,0 +1,65 @@
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
_IS_PHI_4_REGISTERED = False
_IS_PHI_4_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}"
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
# Phi Model Meta
PhiMeta4 = 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
PhiInstructMeta4 = 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_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):
register_phi_4_models(include_original_model=include_original_model)
register_phi_4_instruct_models(include_original_model=include_original_model)
if __name__ == "__main__":
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:
print(f"\u2718 {model_id}")
else:
print(f"\u2713 {model_id}")

117
unsloth/registry/_qwen.py Normal file
View file

@ -0,0 +1,117 @@
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
_IS_QWEN_2_5_REGISTERED = False
_IS_QWEN_2_5_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):
key = f"{base_name}{version}-{size}B"
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):
key = f"{base_name}{version}-VL-{size}B"
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):
key = f"{base_name}-{size}B"
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):
key = f"{base_name}-{size}B-Preview"
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
# Qwen2.5 Model Meta
Qwen_2_5_Meta = 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],
)
# Qwen2.5 VL Model Meta
Qwen_2_5_VLMeta = 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],
)
# 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],
)
# 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_2_5_models(include_original_model: bool = False):
global _IS_QWEN_2_5_REGISTERED
if _IS_QWEN_2_5_REGISTERED:
return
_register_models(Qwen_2_5_Meta, include_original_model=include_original_model)
_IS_QWEN_2_5_REGISTERED = True
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(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
if _IS_QWEN_QWQ_REGISTERED:
return
_register_models(QwenQwQMeta, include_original_model=include_original_model)
_register_models(QwenQVQPreviewMeta, include_original_model=include_original_model)
_IS_QWEN_QWQ_REGISTERED = 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__":
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:
print(f"\u2718 {model_id}")
else:
print(f"\u2713 {model_id}")

View file

@ -0,0 +1,185 @@
import warnings
from dataclasses import dataclass, field
from enum import Enum
class QuantType(Enum):
BNB = "bnb"
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
@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: QuantType = None
description: str = 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: QuantType = None
):
if quant_type != QuantType.NONE:
key = "-".join([key, QUANT_TAG_MAP[quant_type]])
return key
@classmethod
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(
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[QuantType] | dict[str, list[QuantType]] = 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: QuantType = 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, current keys: {MODEL_REGISTRY.keys()}")
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
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):
warnings.warn(f"{model_id} not found on Hugging Face")
model_info = None
else:
raise e
return model_info
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
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:
# 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:
# 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,
org=_org,
base_name=base_name,
version=model_version,
size=size,
instruct_tag=instruct_tag,
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,
)

View file

78
unsloth/utils/hf_hub.py Normal file
View file

@ -0,0 +1,78 @@
from huggingface_hub import HfApi, ModelInfo
_HFAPI: HfApi = None
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.
"""
global _HFAPI
if _HFAPI is None:
_HFAPI = HfApi()
try:
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 list_models(
properties: list[str] = None,
full: bool = False,
sort: str = "downloads",
author: str = "unsloth",
search: str = None,
limit: int = 10,
) -> list[ModelInfo]:
"""
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.
sort: str = The sort order.
author: str = The author of the model.
search: str = The search query for filtering models.
"""
global _HFAPI
if _HFAPI is None:
_HFAPI = HfApi()
if full:
properties = None
models: list[ModelInfo] = _HFAPI.list_models(
author=author,
search=search,
sort=sort,
limit=limit,
expand=properties,
full=full,
)
return models