feat(registry): add public architecture registration and entry-point discovery

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.
+
+<Note>
+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.
+
+</Note>
 
 ### 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.

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()])

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
@@ -458,9 +459,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:
@@ -490,6 +497,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):
@@ -539,4 +559,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()

tests/unit_tests/_transformers/test_registry.py

@@ -596,3 +596,42 @@ 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
+
+    register_architecture("PublicArch", PublicModel)
+    assert inst.get_model_cls_from_model_arch("PublicArch") is PublicModel
+
+
+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