Model registry

[jeromeku]

Model Registry

Create a model registry to systematically track supported models.

Uses:

• testing infrastructure: ensure comprehensive coverage of all model types
• programmatically register, upload, track models on huggingface hub across various permutations (model version, size, base / instruct, quant type, etc.)
• analyze model stats ("downloads", "likes", etc.)

Code Structure

unsloth/model_registry

Models are registered in model_registry.py by providing:
_ModelName_INFO, a dict with the following properties:

• org
• base_name
• instruct_tags
• model_versions
• model_sizes
• model_info_cls

where model_info_cls is a subclass of type:

@dataclass
class ModelInfo:
    org: str
    base_name: str
    version: str
    size: int
    name: str = None  # full model name, constructed from base_name, version, and size unless provided
    is_multimodal: bool = False
    instruct_tag: str = None
    quant_type: Literal["None", "bnb", "unsloth"] = None
    ...

Models are then registered as ModelInfo objects under the global MODEL_REGISTRY dict.

Tests

tests/test_model_registry.py

Tests that models currently registered within model_registry are valid huggingface model repos.

pytest -sv tests/test_model_registry.py

Note that Llama Vision tests will fail since there are the 90b version is missing certain quant types.

TODOs:

Models

• Add llama 3.2 text-only models
• Add llama 3.2 vision models
• Add gemma 3 models
• Add qwen2.5 models
• Add qwen2.5 VL models
• Add phi4 models
• Add deepseek v3
• Add deepseed r1
• Add qwen QwQ

Quant types / model formats

• gguf

HF Hub

• Utils for uploading, tracking, analyzing model usage stats
• Interface to ensure MODEL_REGISTRY stays synced with hf hub

[jeromeku]

Closed in favor of #2255