nemo_automodel.shared.model_utils | NVIDIA NeMo AutoModel

Shared structural model traversal utilities.

Module Contents

Functions

Name	Description
`iter_transformer_and_mtp_blocks`	Yield transformer and MTP blocks without depending on a recipe or component.

Data

_TEXT_MODULE_ATTRS

API

nemo_automodel.shared.model_utils.iter_transformer_and_mtp_blocks(
    model: torch.nn.Module
) -> collections.abc.Iterator[tuple[torch.nn.Module, str, torch.nn.Module]]

Yield transformer and MTP blocks without depending on a recipe or component.

Parameters:

model

nn.Module

Model root containing a transformer layer collection and optional multi-token-prediction layers.

nemo_automodel.shared.model_utils._TEXT_MODULE_ATTRS = ('language_model', 'text_model', 'text_decoder')