From 970b0a9276b0b0717dd846dada4fb65b04e02da2 Mon Sep 17 00:00:00 2001 From: jeromeku Date: Mon, 31 Mar 2025 18:24:15 -0700 Subject: [PATCH] clear registry when executing individual model registration file --- tests/test_model_registry.py | 5 ++-- unsloth/registry/REGISTRY.md | 50 ++++++++++++++++++++++++++++++++++- unsloth/registry/__init__.py | 2 +- unsloth/registry/_deepseek.py | 4 +++ unsloth/registry/_gemma.py | 5 +++- unsloth/registry/_llama.py | 4 ++- unsloth/registry/_mistral.py | 5 +++- unsloth/registry/_phi.py | 5 +++- unsloth/registry/_qwen.py | 5 +++- 9 files changed, 76 insertions(+), 9 deletions(-) diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 3d570af230..f59f4f0dab 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -13,7 +13,7 @@ from dataclasses import dataclass import pytest from huggingface_hub import ModelInfo as HfModelInfo -from unsloth.registry import get_model_info, register_models +from unsloth.registry import register_models, search_models from unsloth.registry._deepseek import register_deepseek_models from unsloth.registry._gemma import register_gemma_models from unsloth.registry._llama import register_llama_models @@ -21,6 +21,7 @@ from unsloth.registry._mistral import register_mistral_models from unsloth.registry._phi import register_phi_models from unsloth.registry._qwen import register_qwen_models from unsloth.registry.registry import MODEL_REGISTRY, QUANT_TAG_MAP, QuantType +from unsloth.utils.hf_hub import get_model_info MODEL_NAMES = [ "llama", @@ -84,7 +85,7 @@ def test_all_model_registration(): def test_quant_type(): # Test that the quant_type is correctly set for model paths # NOTE: for models registered under org="unsloth" with QuantType.NONE aliases QuantType.UNSLOTH - dynamic_quant_models = get_model_info(quant_types=[QuantType.UNSLOTH]) + dynamic_quant_models = search_models(quant_types=[QuantType.UNSLOTH]) assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models) quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH] assert all(quant_tag in m.model_path for m in dynamic_quant_models) \ No newline at end of file diff --git a/unsloth/registry/REGISTRY.md b/unsloth/registry/REGISTRY.md index b794d26be6..8240d686e6 100644 --- a/unsloth/registry/REGISTRY.md +++ b/unsloth/registry/REGISTRY.md @@ -1,6 +1,16 @@ ## Model Registry ### Structure +``` +unsloth + -registry + __init__.py + registry.py + _llama.py + _mistral.py + _phi.py + ... +``` Each model is registered in a separate file within the `registry` module (e.g. `registry/_llama.py`). @@ -41,5 +51,43 @@ class LlamaModelInfo(ModelInfo): return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key) ``` -Once these constructs are defined, the model is registered in the `registry` module by calling `register_models` with the `ModelMeta` and `ModelInfo` classes. +Once these constructs are defined, the model is registered by writing a register_xx_models function. +```python +def register_llama_3_1_models(include_original_model: bool = False): + global _IS_LLAMA_3_1_REGISTERED + if _IS_LLAMA_3_1_REGISTERED: + return + _register_models(LlamaMeta_3_1, include_original_model=include_original_model) + _IS_LLAMA_3_1_REGISTERED = True +``` + +`_register_models` is a helper function that registers the model with the registry. The global `_IS_XX_REGISTERED` is used to prevent duplicate registration. + +Once a model is registered, registry.registry.MODEL_REGISTRY is updated with the model info and can be searched with `registry.search_models`. + +### Tests + +The `tests/test_model_registry.py` file contains tests for the model registry. + +Also, each model registration file is an executable module that checks that all registered models are available on `huggingface_hub`. +```python +python unsloth.registry._llama.py +``` + +Prints the following (abridged) output: +```bash +✓ unsloth/Llama-3.1-8B +✓ unsloth/Llama-3.1-8B-bnb-4bit +✓ unsloth/Llama-3.1-8B-unsloth-bnb-4bit +✓ meta-llama/Llama-3.1-8B +✓ unsloth/Llama-3.1-8B-Instruct +✓ unsloth/Llama-3.1-8B-Instruct-bnb-4bit +✓ unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit +✓ meta-llama/Llama-3.1-8B-Instruct +✓ unsloth/Llama-3.2-1B +✓ unsloth/Llama-3.2-1B-bnb-4bit +✓ unsloth/Llama-3.2-1B-unsloth-bnb-4bit +✓ meta-llama/Llama-3.2-1B +... +``` diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py index a46ab773d8..5874743694 100644 --- a/unsloth/registry/__init__.py +++ b/unsloth/registry/__init__.py @@ -22,7 +22,7 @@ def register_models(): _ARE_MODELS_REGISTERED = True -def get_model_info(org: str = None, base_name: str = None, version: str = None, size: str = None, quant_types: list[QuantType] = None, search_pattern: str = None) -> list[ModelInfo]: +def search_models(org: str = None, base_name: str = None, version: str = None, size: str = None, quant_types: list[QuantType] = None, search_pattern: str = None) -> list[ModelInfo]: """ Get model info from the registry. diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 854a62c00b..153a0e508e 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -163,6 +163,10 @@ register_deepseek_models(include_original_model=True) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_deepseek_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_gemma.py b/unsloth/registry/_gemma.py index 8c47e7e69d..9490c84f2f 100644 --- a/unsloth/registry/_gemma.py +++ b/unsloth/registry/_gemma.py @@ -53,8 +53,11 @@ def register_gemma_models(include_original_model: bool = False): if __name__ == "__main__": - register_gemma_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_gemma_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_llama.py b/unsloth/registry/_llama.py index ec6e39a86d..1c2dd5bf18 100644 --- a/unsloth/registry/_llama.py +++ b/unsloth/registry/_llama.py @@ -100,8 +100,10 @@ def register_llama_models(include_original_model: bool = False): register_llama_3_2_vision_models(include_original_model=include_original_model) if __name__ == "__main__": - register_llama_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_llama_models(include_original_model=True) for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) diff --git a/unsloth/registry/_mistral.py b/unsloth/registry/_mistral.py index c41b1f55b6..44cd1e7646 100644 --- a/unsloth/registry/_mistral.py +++ b/unsloth/registry/_mistral.py @@ -57,8 +57,11 @@ def register_mistral_models(include_original_model: bool = False): register_mistral_small_models(include_original_model=include_original_model) if __name__ == "__main__": - register_mistral_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_mistral_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_phi.py b/unsloth/registry/_phi.py index 9f23c494d5..d06ec8d377 100644 --- a/unsloth/registry/_phi.py +++ b/unsloth/registry/_phi.py @@ -52,8 +52,11 @@ def register_phi_models(include_original_model: bool = False): register_phi_4_instruct_models(include_original_model=include_original_model) if __name__ == "__main__": - register_phi_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_phi_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py index c364f9b099..4417515a77 100644 --- a/unsloth/registry/_qwen.py +++ b/unsloth/registry/_qwen.py @@ -104,8 +104,11 @@ def register_qwen_models(include_original_model: bool = False): register_qwen_qwq_models(include_original_model=include_original_model) if __name__ == "__main__": - register_qwen_models(include_original_model=True) from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info + MODEL_REGISTRY.clear() + + register_qwen_models(include_original_model=True) + for model_id, model_info in MODEL_REGISTRY.items(): model_info = _check_model_info(model_id) if model_info is None: