Skip to content

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

Open
pstjohn wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
pstjohn:u-2-public-registry
Open

feat(registry): add public architecture registration and entry-point discovery#3645
pstjohn wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
pstjohn:u-2-public-registry

Conversation

@pstjohn

@pstjohn pstjohn commented Aug 24, 2026

Copy link
Copy Markdown

Adds public out-of-tree architecture registration.

Custom model architectures (e.g. multimodal VLMs built by decorating a Hub-resolved text-only model) must register with AutoModel's model registry before construction. Today that requires importing the private nemo_automodel._transformers.registry.ModelRegistry and calling register() in every process that resolves architectures — including isolated Ray workers and CLI subprocesses where import-order dependencies are fragile.

_ModelRegistry now scans nemo_automodel.architectures entry points on first access, same shape as vLLM's vllm.general_plugins. Entries are arch_name = module.path:ClassName and are added to the lazy mapping without importing the target module; import happens on first architecture resolution, exactly like the built-in MODEL_ARCH_MAPPING entries.

We also export a public register_architecture():

from nemo_automodel import register_architecture
register_architecture("MyArch", MyModelClass)

This just delegates to the same _ModelRegistry.register, so explicit and entry-point-discovered registrations share one code path.

With entry-point discovery, you can now declare their architecture once in your pyproject.toml and every process discovers it automatically:

[project.entry-points."nemo_automodel.architectures"]
MyArch = "my_package.model:MyModelClass"

This uses Python's standard importlib.metadata.entry_points() (PEP 621).

@github-actions

github-actions Bot commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

@HuiyingLi

Copy link
Copy Markdown
Contributor

Wow thank you @pstjohn for adding it!

@jgerh jgerh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Completed tech pubs review of docs/model-coverage/overview.mdx and provided some copyedits and suggested text revisions.

Comment thread docs/model-coverage/overview.mdx Outdated
Comment thread docs/model-coverage/overview.mdx Outdated
Comment thread docs/model-coverage/overview.mdx Outdated
Comment thread docs/model-coverage/overview.mdx Outdated
Comment thread docs/model-coverage/overview.mdx Outdated
Comment thread docs/model-coverage/overview.mdx Outdated
Comment thread docs/model-coverage/overview.mdx Outdated
Co-authored-by: jgerh <163925524+jgerh@users.noreply.github.com>
Signed-off-by: Peter St. John <pstjohn@nvidia.com>

@jgerh jgerh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Completed tech pubs review of docs/model-coverage/overview.mdx and provided copyedits and duplicate text deletions.

Comment thread docs/model-coverage/overview.mdx Outdated
Comment thread docs/model-coverage/overview.mdx
Comment thread docs/model-coverage/overview.mdx Outdated
Co-authored-by: jgerh <163925524+jgerh@users.noreply.github.com>
Signed-off-by: Peter St. John <pstjohn@nvidia.com>
@pstjohn

pstjohn commented Aug 27, 2026

Copy link
Copy Markdown
Author

/ok to test ae00137

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants