nemo_automodel.components.models | NVIDIA NeMo AutoModel

Convenience model builders for NeMo Automodel.

Currently includes: • build_gpt2_model – returns a GPT-2 causal language model (Flash-Attention-2 by default).

Subpackages

Submodules

Package Contents

Classes

Name	Description
`_MissingModelFinder`	Produces a helpful error when importing a non-existent model subpackage.

Functions

Name	Description
`__getattr__`	-
`_available_model_submodules`	Return the set of model sub-package names shipped with this installation.
`_make_upgrade_message`	-
`build_gpt2_model`	Instantiate and return a pure-PyTorch GPT-2 language model.

Data

_MODELS_DIR

_PACKAGE_PREFIX

API

class nemo_automodel.components.models._MissingModelFinder()

Bases: MetaPathFinder

Produces a helpful error when importing a non-existent model subpackage.

Installed at the end of sys.meta_path so it is only consulted after all real finders have already returned None. For any import of the form nemo_automodel.components.models.<name> (direct child only), it raises a ModuleNotFoundError with upgrade instructions instead of the default unhelpful message.

nemo_automodel.components.models._MissingModelFinder.find_spec(
    fullname,
    path,
    target = None
)

nemo_automodel.components.models.__getattr__(
    name: str
)

nemo_automodel.components.models._available_model_submodules() -> set[str]

Return the set of model sub-package names shipped with this installation.

nemo_automodel.components.models._make_upgrade_message(
    name: str
) -> str

nemo_automodel.components.models.gpt2.build_gpt2_model(
    vocab_size: int = 50257,
    n_positions: int = 2048,
    n_ctx: int | None = None,
    n_embd: int = 768,
    n_layer: int = 12,
    n_head: int = 12,
    bos_token_id: int = 50256,
    eos_token_id: int = 50256,
    attn_implementation: str = 'flash_attention_2',
    extra_cfg: typing.Any = {}
) -> torch.nn.Module

Instantiate and return a pure-PyTorch GPT-2 language model.

The function intentionally keeps the same signature as the original wrapper so existing YAML/CLI configurations continue to work. Extra keyword arguments are quietly ignored.

nemo_automodel.components.models._MODELS_DIR = pathlib.Path(__file__).parent

nemo_automodel.components.models._PACKAGE_PREFIX = __name__ + '.'