Merge pull request #2255 from jeromeku/registry-refactor
Registry refactor
This commit is contained in:
commit
fd20192aef
15 changed files with 1306 additions and 0 deletions
177
.gitignore
vendored
Normal file
177
.gitignore
vendored
Normal 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
|
||||
|
|
@ -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
0
tests/__init__.py
Normal file
91
tests/test_model_registry.py
Normal file
91
tests/test_model_registry.py
Normal 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)
|
||||
110
unsloth/registry/REGISTRY.md
Normal file
110
unsloth/registry/REGISTRY.md
Normal 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
|
||||
51
unsloth/registry/__init__.py
Normal file
51
unsloth/registry/__init__.py
Normal 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
|
||||
179
unsloth/registry/_deepseek.py
Normal file
179
unsloth/registry/_deepseek.py
Normal 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)
|
||||
66
unsloth/registry/_gemma.py
Normal file
66
unsloth/registry/_gemma.py
Normal 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
113
unsloth/registry/_llama.py
Normal 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}")
|
||||
70
unsloth/registry/_mistral.py
Normal file
70
unsloth/registry/_mistral.py
Normal 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
65
unsloth/registry/_phi.py
Normal 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
117
unsloth/registry/_qwen.py
Normal 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}")
|
||||
185
unsloth/registry/registry.py
Normal file
185
unsloth/registry/registry.py
Normal 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,
|
||||
)
|
||||
0
unsloth/utils/__init__.py
Normal file
0
unsloth/utils/__init__.py
Normal file
78
unsloth/utils/hf_hub.py
Normal file
78
unsloth/utils/hf_hub.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue