`nemo_automodel._transformers.model_capabilities`#

Programmatic query API for NeMo AutoModel’s per-architecture capability registry.

Two declaration patterns are supported on registered model classes:

Static – nested ModelCapabilities dataclass on the class. Used by classes whose capability profile does not depend on config.
Dynamic – get_capabilities(cls, config) classmethod. Used by classes that serve multiple checkpoint variants (e.g. Gemma4 MoE vs. dense vs. audio).

A class must declare capabilities via exactly one of these patterns.

func:: query_capabilities dispatches to whichever is present and always returns the canonical :class:ModelCapabilities defined here.

Module Contents#

Classes#

ModelCapabilities

Canonical parallelism capability flags for a model architecture.

Functions#

`_to_canonical`	Re-pack any object with the four `supports_*` fields into the canonical type.
`_resolve_class_from_arch`	Look up a registered model class by architecture name. Raises on miss.
`_arch_from_config`	Extract the primary architecture name from an HF `PretrainedConfig`.
`_dispatch`	Apply the static/dynamic capability rules for a resolved class.
`query_capabilities`	Resolve declared parallelism capabilities for a model.

Data#

__all__

API#

nemo_automodel._transformers.model_capabilities.__all__#: [‘ModelCapabilities’, ‘query_capabilities’]

class nemo_automodel._transformers.model_capabilities.ModelCapabilities#

Canonical parallelism capability flags for a model architecture.

All fields are conservative defaults; a flag is True only when the corresponding feature has a working, verified implementation for that architecture (or variant, in the dynamic-dispatch case).

.. attribute:: supports_tp

Tensor parallelism.

.. attribute:: supports_cp

Context parallelism.

.. attribute:: supports_pp

Pipeline parallelism.

.. attribute:: supports_ep

Expert parallelism (MoE).

supports_tp: bool#: False

supports_cp: bool#: False

supports_pp: bool#: False

supports_ep: bool#: False

nemo_automodel._transformers.model_capabilities._to_canonical( caps_obj, ) → nemo_automodel._transformers.model_capabilities.ModelCapabilities#

Re-pack any object with the four supports_* fields into the canonical type.

Per-class nested ModelCapabilities dataclasses are their own types; this converts them to the canonical :class:ModelCapabilities so callers see a single, stable type.

nemo_automodel._transformers.model_capabilities._resolve_class_from_arch(arch: str)#: Look up a registered model class by architecture name. Raises on miss.

nemo_automodel._transformers.model_capabilities._arch_from_config(config) → str#: Extract the primary architecture name from an HF PretrainedConfig.

nemo_automodel._transformers.model_capabilities._dispatch( model_cls, config, ) → nemo_automodel._transformers.model_capabilities.ModelCapabilities#: Apply the static/dynamic capability rules for a resolved class.

nemo_automodel._transformers.model_capabilities.query_capabilities( target: Union[str, transformers.PretrainedConfig, torch.nn.Module, type], *, trust_remote_code: bool = False, ) → nemo_automodel._transformers.model_capabilities.ModelCapabilities#

Resolve declared parallelism capabilities for a model.

Accepted input forms:

HF model id (str) – AutoConfig.from_pretrained is used to load the config (no weights), the architecture is read from the config, and the registered NeMo class is looked up via :data:ModelRegistry.
HF config object – the architecture is read from the config.
Model instance (nn.Module) – type(target) resolves the class and target.config provides the variant info needed by dynamic dispatch.
Registered model class (type) – usable only for classes with a static nested ModelCapabilities dataclass. Classes with dynamic dispatch raise :class:ValueError because no config is available.

Parameters:

target – HF model id, PretrainedConfig, model instance, or registered NeMo model class.
trust_remote_code – Forwarded to AutoConfig.from_pretrained when target is a string. Ignored otherwise.

Returns:

A :class:ModelCapabilities (canonical type, identical across all inputs and all model classes).

Raises:

TypeError – If target is not a supported input form.
KeyError – If the resolved architecture is not registered.
ValueError – If a dynamic-dispatch class is queried via the bare class (no config available) or a config has no architectures field.
AttributeError – If the resolved class declares no capabilities.

nemo_automodel._transformers.model_capabilities#