diff --git a/docs/model-coverage/overview.mdx b/docs/model-coverage/overview.mdx index 7e590cfc96..a57cc74083 100644 --- a/docs/model-coverage/overview.mdx +++ b/docs/model-coverage/overview.mdx @@ -1,6 +1,6 @@ --- title: "Model Coverage Overview" -description: "" +description: "Hugging Face Auto class support, release history, day-0 coverage, and custom architecture registration for NeMo AutoModel." position: 1 --- NeMo AutoModel integrates with Hugging Face `transformers`. Any LLM or VLM that can be instantiated through `transformers` can also be used with NeMo AutoModel, subject to runtime, third-party software dependencies, and feature compatibility. @@ -14,7 +14,7 @@ NeMo AutoModel integrates with Hugging Face `transformers`. Any LLM or VLM that | Block-Diffusion LLMs | Text Generation (Diffusion LLM) | Supported | See [Diffusion LLM model list](/model-coverage/dllm/overview) | | `AutoModelForImageTextToText` | Image-Text-to-Text (VLM) | Supported | See [VLM model list](/model-coverage/vision-language-models/overview) | | Custom multimodal models | Unified multimodal training | Supported | See [Multimodal model list](/model-coverage/multimodal/overview) | -| `AutoModelForSequenceClassification` | Sequence Classification | Work in Progress | Early support. Interfaces might change | +| `AutoModelForSequenceClassification` | Sequence Classification | Work in Progress | Early support; interfaces might change | | Diffusers Pipelines | Diffusion Generation (T2I, T2V) | Supported | See [Diffusion model list](/model-coverage/diffusion/overview) | | `NeMoAutoModelBiEncoder` | Embedding Models | Supported | See [Embedding model list](/model-coverage/embedding-models/overview) | | `NeMoAutoModelCrossEncoder` | Reranking Models | Supported | See [Reranking model list](/model-coverage/reranking-models/overview) | @@ -25,15 +25,15 @@ The table below tracks when model support and key features were added across NeM | Release | Date | New Models | Key Features | |---------|------|------------|--------------| -| **0.3.0** (upcoming) | Not announced | Kimi-VL, Kimi-K25-VL, Gemma 3n, Nemotron-Parse, Qwen3-VL-MoE, Qwen3-Omni, InternVL 3.5, Ministral3, Phi-4-multimodal, Devstral-Small-2, Step-3.5-Flash, Qwen3-Next, Nemotron-3-Nano-30B, FLUX.1-dev, Wan 2.1 T2V, HunyuanVideo 1.5 | MoE LoRA, expanded VLM coverage, diffusion model training (flow matching) | -| **0.2.0** | Dec 2025 | GPT-OSS 20B/120B, Qwen3, Qwen3-MoE, GLM-4/4-MoE, Qwen2.5-VL, Qwen3-VL | Single- and multi-turn tool calling, streaming dataset, QAT for SFT, sequence classification, async DCP checkpointing, MLflow, CP and sequence packing for MoE | -| **0.1.0** | Oct 2025 | DeepSeek V3/V3.2, more than 40 LLM architectures, Gemma 3 VLM | Pretraining, knowledge distillation, FP8 (torchao), pipeline parallelism, HSDP, auto pipelining, ColumnMapped dataset | -| **0.1.0a0** | Sep 2025 | Initial LLM and VLM support (Llama, Mistral, Qwen2, Gemma, Phi, and more) | MegatronFSDP, packed sequences, Triton LoRA kernels | +| **0.3.0** (upcoming) | Not announced | Kimi-VL, Kimi-K25-VL, Gemma 3n, Nemotron-Parse, Qwen3-VL-MoE, Qwen3-Omni, InternVL 3.5, Ministral3, Phi-4-multimodal, Devstral-Small-2, Step-3.5-Flash, Qwen3-Next, Nemotron-3-Nano-30B, FLUX.1-dev, Wan 2.1 T2V, HunyuanVideo 1.5 | Mixture-of-Experts (MoE) LoRA, expanded VLM coverage, diffusion model training (flow matching) | +| **0.2.0** | December 2025 | GPT-OSS 20B/120B, Qwen3, Qwen3-MoE, GLM-4/4-MoE, Qwen2.5-VL, Qwen3-VL | Single- and multi-turn tool calling, streaming dataset, QAT for SFT, sequence classification, async DCP checkpointing, MLflow, CP and sequence packing for MoE | +| **0.1.0** | October 2025 | DeepSeek V3/V3.2, more than 40 LLM architectures, Gemma 3 VLM | Pretraining, knowledge distillation, FP8 (torchao), pipeline parallelism, HSDP, auto pipelining, ColumnMapped dataset | +| **0.1.0a0** | September 2025 | Initial LLM and VLM support (Llama, Mistral, Qwen2, Gemma, Phi, and more) | MegatronFSDP, packed sequences, Triton LoRA kernels | ## Day-0 Support - NeMo AutoModel closely tracks the latest `transformers` version and updates its dependency regularly. -- New models released on the Hugging Face Hub may require the latest `transformers` version, necessitating a package upgrade. +- New models released on the Hugging Face Hub might require the latest `transformers` version, necessitating a package upgrade. - The team is developing a CI pipeline that automatically updates the supported `transformers` version when a new release is detected, enabling faster day-0 support. @@ -42,17 +42,63 @@ The table below tracks when model support and key features were added across NeM NeMo AutoModel includes a custom model registry that allows teams to: - Add custom implementations to extend support to models not yet covered upstream. -- Provide optimized or faster implementations for specific models while retaining the same AutoModel interface. +- Provide optimized or faster implementations for specific models while retaining the same NeMo AutoModel interface. + +### Register an Architecture + +The registry matches an architecture name against the first value in the checkpoint's `config.json` +`architectures` list. The name is case-sensitive. The registered class must be a `torch.nn.Module` class that is +compatible with the selected `NeMoAutoModel*` loader and accepts the resolved Hugging Face config as its first +constructor argument. + +#### Register in Python + +Call `register_architecture` before constructing or loading the model: + +```python +from nemo_automodel import NeMoAutoModelForCausalLM, register_architecture + +from my_package.models import MyModelForCausalLM + +register_architecture("MyModelForCausalLM", MyModelForCausalLM) +model = NeMoAutoModelForCausalLM.from_pretrained("my-org/my-model") +``` + +Registering a built-in or previously registered name raises `ValueError`, even when the same class is registered +again. Pass `exist_ok=True` only when you intentionally want to replace the existing model class. + +#### Register from an Installed Package + +An installed package can advertise model classes without requiring application startup code. Add an entry point for +each architecture to the package's `pyproject.toml`: + +```toml +[project.entry-points."nemo_automodel.architectures"] +MyModelForCausalLM = "my_package.models:MyModelForCausalLM" +``` + +The entry-point name is the architecture name, and its value must use the `module.path:ClassName` format. NeMo +AutoModel discovers these entry points when its model registry initializes and imports the target module only when it +resolves that architecture. Install the package before starting the Python process. If an entry-point name conflicts +with a built-in or another discovered architecture, NeMo AutoModel skips it and logs a warning; use +`register_architecture(..., exist_ok=True)` in application code for an intentional override. + + +Architecture registration selects a model implementation after the Hugging Face config is resolved. It does not +register a new `model_type`, so the checkpoint config must already be loadable by the installed versions of +Hugging Face `transformers` or NeMo AutoModel. + + ### Ready-to-Run Architectures The following table lists architectures represented by the ready-to-run YAML recipes in this repository. It includes both NeMo-native implementations and models that use the standard Hugging Face implementation path. -This is a practical starting set rather than an exhaustive compatibility list. NeMo AutoModel can also work with additional models supported by the installed version of the Hugging Face `transformers` library, although models without a checked-in recipe may require some configuration for a particular training setup. +This is a practical starting set rather than an exhaustive compatibility list. NeMo AutoModel can also work with additional models supported by the installed version of the Hugging Face `transformers` library, although models without a checked-in recipe might require some configuration for a particular training setup. {/* BEGIN GENERATED MODEL ARCHITECTURES */} {/* END GENERATED MODEL ARCHITECTURES */} ## Having Issues? -If a model from the Hub does not work as expected, see [Troubleshooting](/model-coverage/troubleshooting) for common issues and solutions. +If a model from the Hugging Face Hub does not work as expected, see [Troubleshooting](/model-coverage/troubleshooting) for common issues and solutions. diff --git a/nemo_automodel/__init__.py b/nemo_automodel/__init__.py index 740728f570..9b45670f78 100644 --- a/nemo_automodel/__init__.py +++ b/nemo_automodel/__init__.py @@ -59,6 +59,7 @@ "NeMoAutoDiffusionPipeline": ("nemo_automodel._diffusers.auto_diffusion_pipeline", "NeMoAutoDiffusionPipeline"), "ModelCapabilities": ("nemo_automodel._transformers.model_capabilities", "ModelCapabilities"), "query_capabilities": ("nemo_automodel._transformers.model_capabilities", "query_capabilities"), + "register_architecture": ("nemo_automodel._transformers.registry", "register_architecture"), } __all__ = sorted([*_SUBMODULES, "__version__", "__package_name__", *_LAZY_ATTRS.keys()]) diff --git a/nemo_automodel/_transformers/registry.py b/nemo_automodel/_transformers/registry.py index d99cc86e01..0451537a40 100644 --- a/nemo_automodel/_transformers/registry.py +++ b/nemo_automodel/_transformers/registry.py @@ -14,6 +14,7 @@ import importlib +import importlib.metadata import logging from collections import OrderedDict from dataclasses import dataclass, field @@ -489,9 +490,15 @@ def __setitem__(self, key: str, value: Type[nn.Module]) -> None: self._extra[key] = value def register(self, key: str, value: Type[nn.Module], exist_ok: bool = False) -> None: - """Register a model class under the given architecture name.""" - if not exist_ok and key in self._extra: + """Register a model class under the given architecture name. + + Conflicts, including built-in architectures, require ``exist_ok=True``. + """ + if not exist_ok and (key in self._auto_map or key in self._extra): raise ValueError(f"Duplicated model implementation for {key}") + + self._auto_map.pop(key, None) + self._loaded.pop(key, None) self._extra[key] = value def has_tag(self, key: str, tag: str) -> bool: @@ -521,6 +528,19 @@ def __post_init__(self): if self.model_arch_name_to_cls is None: self.model_arch_name_to_cls = _LazyArchMapping(MODEL_ARCH_MAPPING) self._retrieval_archs = self.model_arch_name_to_cls.keys_with_tag("retrieval") + self._discover_entry_points() + + def _discover_entry_points(self) -> None: + """Register architecture entry points without importing their modules.""" + mapping = self.model_arch_name_to_cls + for ep in importlib.metadata.entry_points(group="nemo_automodel.architectures"): + if ep.name in mapping.keys(): + logger.warning("Architecture %s is already registered; skipping entry point %s", ep.name, ep.value) + continue + module_path, _, class_name = ep.value.rpartition(":") + if not module_path or not class_name: + raise ValueError(f"Entry point {ep.name!r} value must be module.path:ClassName, got {ep.value!r}") + mapping._auto_map[ep.name] = (module_path, class_name) @property def supported_models(self): @@ -570,4 +590,19 @@ def get_registry(): return _ModelRegistry() +def register_architecture(arch_name: str, model_cls: type[nn.Module], *, exist_ok: bool = False) -> None: + """Register a custom model class for an architecture name. + + Args: + arch_name: Architecture name (e.g. ``"LlavaExampleNemotronForCausalLM"``). + model_cls: The model class (not a string path; the class object itself). + exist_ok: If True, replace an existing registration. + + Raises: + ValueError: If *arch_name* is already a built-in or registered to a + different class and *exist_ok* is False. + """ + ModelRegistry.register(arch_name, model_cls, exist_ok=exist_ok) + + ModelRegistry = get_registry() diff --git a/tests/unit_tests/_transformers/test_registry.py b/tests/unit_tests/_transformers/test_registry.py index 546c3a77ed..f4c75b03e0 100644 --- a/tests/unit_tests/_transformers/test_registry.py +++ b/tests/unit_tests/_transformers/test_registry.py @@ -92,6 +92,29 @@ class ReplacementClass: assert inst.model_arch_name_to_cls["MyArch"] is ReplacementClass +def test_builtin_register_requires_explicit_override(): + """Built-in architecture names are preserved unless the caller opts into replacement.""" + from nemo_automodel._transformers.registry import _LazyArchMapping + + class BuiltInClass: + pass + + class ReplacementClass: + pass + + mapping = _LazyArchMapping({"BuiltInArch": ("built.in", "BuiltInClass")}) + mapping._modules["built.in"] = types.SimpleNamespace(BuiltInClass=BuiltInClass) + + with pytest.raises(ValueError, match="Duplicated model implementation for BuiltInArch"): + mapping.register("BuiltInArch", ReplacementClass) + assert mapping["BuiltInArch"] is BuiltInClass + + mapping.register("BuiltInArch", ReplacementClass, exist_ok=True) + assert mapping["BuiltInArch"] is ReplacementClass + assert "BuiltInArch" not in mapping._auto_map + assert "BuiltInArch" not in mapping._loaded + + def test_supported_models_and_getter(): from nemo_automodel._transformers import registry as reg @@ -596,3 +619,83 @@ def test_minimax_m3_vl_config_overrides_transformers_builtin(): ) # The concrete field the crash was about: our vision sub-config must default it. assert MiniMaxM3VLConfig().vision_config.rope_theta is not None + + +def test_public_register_architecture(monkeypatch): + """register_architecture is importable from nemo_automodel.""" + from nemo_automodel import register_architecture + from nemo_automodel._transformers import registry as reg + + inst = _new_registry_instance(reg) + monkeypatch.setattr(reg, "ModelRegistry", inst) + + class PublicModel: + pass + + class ReplacementModel: + pass + + register_architecture("PublicArch", PublicModel) + assert inst.get_model_cls_from_model_arch("PublicArch") is PublicModel + + register_architecture("PublicArch", ReplacementModel, exist_ok=True) + assert inst.get_model_cls_from_model_arch("PublicArch") is ReplacementModel + + +def test_entry_point_discovery_adds_lazy_entry(monkeypatch): + """Entry points in nemo_automodel.architectures are discovered and lazily loaded.""" + import importlib.metadata + + from nemo_automodel._transformers import registry as reg + + class EntryModel: + pass + + fake_ep = types.SimpleNamespace(name="EntryArch", value="fake.module:EntryModel") + monkeypatch.setattr( + importlib.metadata, + "entry_points", + lambda **kwargs: [fake_ep] if kwargs.get("group") == "nemo_automodel.architectures" else [], + ) + + inst = _new_registry_instance(reg) + assert "EntryArch" in inst.model_arch_name_to_cls._auto_map + + # First resolution triggers import. + inst.model_arch_name_to_cls._modules["fake.module"] = types.SimpleNamespace(EntryModel=EntryModel) + assert inst.get_model_cls_from_model_arch("EntryArch") is EntryModel + + +def test_entry_point_discovery_skips_registered_architecture(monkeypatch, caplog): + """Entry points cannot silently replace an architecture already in the registry.""" + import importlib.metadata + + from nemo_automodel._transformers import registry as reg + + fake_ep = types.SimpleNamespace(name="ExistingArch", value="plugin.module:PluginModel") + monkeypatch.setattr(importlib.metadata, "entry_points", lambda **kwargs: [fake_ep]) + + mapping = reg._LazyArchMapping({"ExistingArch": ("built.in", "BuiltInModel")}) + reg._ModelRegistry(model_arch_name_to_cls=mapping) + + assert mapping._auto_map["ExistingArch"] == ("built.in", "BuiltInModel") + assert ( + "Architecture ExistingArch is already registered; skipping entry point plugin.module:PluginModel" in caplog.text + ) + + +@pytest.mark.parametrize("entry_point_value", [":EntryModel", "fake.module:"]) +def test_entry_point_discovery_rejects_invalid_value(monkeypatch, entry_point_value): + """Architecture entry points require both a module path and class name.""" + import importlib.metadata + + from nemo_automodel._transformers import registry as reg + + fake_ep = types.SimpleNamespace(name="BrokenArch", value=entry_point_value) + monkeypatch.setattr(importlib.metadata, "entry_points", lambda **kwargs: [fake_ep]) + + with pytest.raises( + ValueError, + match=r"Entry point 'BrokenArch' value must be module\.path:ClassName", + ): + _new_registry_instance(reg)