nemo_automodel.components.utils.model_utils

Module Contents

Classes

Name	Description
`FreezeConfig`	Typed schema for the `freeze_config` recipe section.
`ModuleSelector`	Select modules by exact canonical path or case-sensitive shell-style glob.

Functions

Name	Description
`_apply_module_selectors`	Set `requires_grad` on every module matched by `selectors`, recursively.
`_freeze_module_by_attribute_and_patterns`	Freeze a legacy model attribute and modules matching name substrings.
`_get_cached_forward_signature`	Best-effort retrieval cached by the underlying `forward` callable.
`_get_forward_signature`	Retrieve `model.forward` once per callable, preserving live patches.
`_get_logical_numel`	Return the logical number of elements for a parameter,
`_get_model_param_stats`	Get the number of trainable parameters and the L2 norm of the model.
`_parse_module_selector`	Parse one `freeze_modules`/`unfreeze_modules` entry into a ModuleSelector.
`_supports_logits_to_keep`	Check if the model supports logits_to_keep.
`_supports_seq_lens`	Check if the model’s forward() accepts seq_lens.
`apply_parameter_freezing`	Apply parameter freezing based on a typed FreezeConfig.
`cast_mixed_dtype_params_to_bf16`	Cast fp32 parameters and buffers to bf16 for FSDP2 compatibility.
`count_model_parameters`	Count total and trainable parameters. Safe to call on meta-device models.
`enable_radio_vit_fused_attn`	Route RADIO ViT attention through `F.scaled_dot_product_attention`.
`filter_forward_kwargs`	Drop kwargs that `model.forward` does not accept.
`freeze_deepseek_v4_indexer_params`	Freeze DeepSeek V4 indexer params that only feed discrete top-k masks.
`freeze_minimax_m3_indexer_params`	Freeze MiniMax M3 lightning-indexer params that only feed discrete top-k masks.
`freeze_unused_kv_sharing_params`	Freeze dead K/V parameters in KV-shared layers.
`get_lm_head_module`	Return the model’s LM head module, if one can be found.
`get_lm_head_weight`	Return the model’s LM-head weight, materializing DTensor weights when needed.
`init_empty_weights`	A context manager under which models are initialized with all parameters on the specified device.
`parse_freeze_config`	Validate raw `freeze_config` YAML data and return a typed FreezeConfig.
`print_trainable_parameters`	Print the number of trainable parameters in the model.
`resolve_trust_remote_code`	Whitelist NVIDIA models to allow remote code execution.
`skip_random_init`	Context manager to skip random weight initialization when loading pretrained models.
`squeeze_input_for_thd`	Squeeze batch dimension and prepare inputs for THD (total, hidden, depth) format.

Data

VLM_INPUT_KEYS

_FREEZE_CONFIG_OPTIONS

logger

API

class nemo_automodel.components.utils.model_utils.FreezeConfig(
    freeze_modules: list[nemo_automodel.components.utils.model_utils.ModuleSelector] | None = None,
    unfreeze_modules: list[nemo_automodel.components.utils.model_utils.ModuleSelector] | None = None,
    freeze_vision_tower: bool | None = None,
    freeze_audio_tower: bool = False,
    freeze_language_model: bool = False,
    freeze_video_embedder: bool = False
)

Dataclass

Typed schema for the freeze_config recipe section.

Trainability semantics, in application order:

Full fine-tuning preserves the model’s existing trainability state; PEFT establishes a LoRA-trainable, base-frozen baseline.
freeze_modules recursively freezes the selected modules.
unfreeze_modules recursively unfreezes the selected modules and wins on overlap.
Framework-required freezes (dead K/V projections, indexer parameters) are applied by the infrastructure afterwards and remain protected.
The final trainable parameter set is validated before the optimizer is constructed.

The legacy modality booleans remain supported. freeze_vision_tower defaults to True only for legacy-only configurations; a configuration that declares freeze_modules or unfreeze_modules uses explicit-selector semantics and does not implicitly freeze vision modules.

freeze_audio_tower

bool = False

freeze_language_model

bool = False

freeze_modules

list[ModuleSelector] | None = None

freeze_video_embedder

bool = False

freeze_vision_tower

bool | None = None

unfreeze_modules

list[ModuleSelector] | None = None

nemo_automodel.components.utils.model_utils.FreezeConfig.__post_init__() -> None

Validate selectors and legacy compatibility options.

nemo_automodel.components.utils.model_utils.FreezeConfig.has_generic_selectors() -> bool

Return whether either generic selector field was explicitly declared.

class nemo_automodel.components.utils.model_utils.ModuleSelector(
    path: str | None = None,
    glob: str | None = None
)

Dataclass

Select modules by exact canonical path or case-sensitive shell-style glob.

Exactly one of path or glob must be set. Both match against the full canonical module path (activation-checkpoint and torch.compile wrapper components stripped): path is an exact match, while glob uses :func:fnmatch.fnmatchcase semantics where * also crosses . separators. Matching is recursive: a selected module’s entire subtree is frozen or unfrozen.

glob

str | None = None

path

str | None = None

nemo_automodel.components.utils.model_utils.ModuleSelector.__post_init__() -> None

Validate that exactly one matching mode is configured.

nemo_automodel.components.utils.model_utils.ModuleSelector.describe() -> str

Return the selector in its key: value configuration form.

nemo_automodel.components.utils.model_utils.ModuleSelector.matches(
    module_path: str
) -> bool

Return whether this selector matches a canonical module path.

Parameters:

module_path

str

Canonical fully qualified module path.

Returns: bool

Whether the selector matches module_path.

nemo_automodel.components.utils.model_utils._apply_module_selectors(
    model: torch.nn.Module,
    selectors: list[nemo_automodel.components.utils.model_utils.ModuleSelector],
    requires_grad: bool,
    strict: bool,
    field_name: str
) -> None

Set requires_grad on every module matched by selectors, recursively.

Module paths are canonicalized (activation-checkpoint and _orig_mod wrapper components stripped) so selectors keep matching after model surgery that wraps or renames parameter-holding modules.

Parameters:

model

nn.Module

The model (or pipeline-parallel model part) to modify.

selectors

list[ModuleSelector]

Typed selectors to resolve against the module hierarchy.

requires_grad

bool

Trainability to apply to matched modules.

strict

bool

When True, raise if a selector matches no parameters. Rebinding after parallelization uses False because sharding may relocate or regroup the selected modules (e.g. pipeline stages hold only a part of the model).

field_name

str

Owning FreezeConfig field name, used in error messages.

Raises:

ValueError: If strict and a selector matches no parameters.

nemo_automodel.components.utils.model_utils._freeze_module_by_attribute_and_patterns(
    model,
    attribute_name,
    name_patterns
)

Freeze a legacy model attribute and modules matching name substrings.

nemo_automodel.components.utils.model_utils._get_cached_forward_signature(
    forward_callable: typing.Callable[..., typing.Any]
) -> inspect.Signature | None

Best-effort retrieval cached by the underlying forward callable.

nemo_automodel.components.utils.model_utils._get_forward_signature(
    model: torch.nn.Module
) -> inspect.Signature | None

Retrieve model.forward once per callable, preserving live patches.

nemo_automodel.components.utils.model_utils._get_logical_numel(
    param
) -> int

Return the logical number of elements for a parameter, accounting for quantized (packed) storage.

For bitsandbytes 4-bit params (Params4bit), the physical tensor packs multiple values per byte. We recover the logical count from the original shape stored in param.quant_state.

nemo_automodel.components.utils.model_utils._get_model_param_stats(
    model: torch.nn.Module
) -> tuple[int, int, float]

Get the number of trainable parameters and the L2 norm of the model.

Parameters:

model

nn.Module

Model to analyze

Returns: int

int

nemo_automodel.components.utils.model_utils._parse_module_selector(
    entry: typing.Any,
    field_name: str
) -> nemo_automodel.components.utils.model_utils.ModuleSelector

Parse one freeze_modules/unfreeze_modules entry into a ModuleSelector.

Parameters:

entry

Any

Mapping with exactly one of path or glob, or an existing ModuleSelector.

field_name

str

Owning list field name, used in error messages.

Returns: ModuleSelector

The validated ModuleSelector.

Raises:

ValueError: If the entry is not a {path: ...} / {glob: ...} mapping or contains unknown keys.

nemo_automodel.components.utils.model_utils._supports_logits_to_keep(
    model: torch.nn.Module
) -> bool

Check if the model supports logits_to_keep.

Parameters:

model

nn.Module

The model to check.

Returns: bool

True if the model supports logits_to_keep, False otherwise.

nemo_automodel.components.utils.model_utils._supports_seq_lens(
    model: torch.nn.Module
) -> bool

Check if the model’s forward() accepts seq_lens.

Returns True if:

forward() has an explicit seq_lens parameter, OR
forward() has **kwargs (so it won’t crash if seq_lens is passed)

Returns False otherwise (passing seq_lens would cause “unexpected kwarg” error).

nemo_automodel.components.utils.model_utils.apply_parameter_freezing(
    model: torch.nn.Module,
    freeze_config: nemo_automodel.components.utils.model_utils.FreezeConfig | collections.abc.Mapping[str, typing.Any],
    strict: bool = True
) -> None

Apply parameter freezing based on a typed FreezeConfig.

Application order: legacy modality booleans and freeze_modules freeze, then unfreeze_modules unfreezes and wins on overlap. Framework-required freezes (dead K/V projections, indexer parameters) are applied separately by the infrastructure after this function, before and after sharding.

Parameters:

model

nn.Module

The model to apply freezing to.

freeze_config

FreezeConfig | Mapping[str, Any]

Typed freeze configuration or a raw mapping retained for compatibility with direct callers. Raw mappings are validated and converted to FreezeConfig before use.

strict

boolDefaults to True

When True, raise if a freeze_modules/unfreeze_modules selector matches no parameters. Set False when rebinding the policy onto a post-parallelization model part.

nemo_automodel.components.utils.model_utils.cast_mixed_dtype_params_to_bf16(
    model
)

Cast fp32 parameters and buffers to bf16 for FSDP2 compatibility.

nemo_automodel.components.utils.model_utils.count_model_parameters(
    model: torch.nn.Module
) -> tuple[int, int]

Count total and trainable parameters. Safe to call on meta-device models.

Parameters:

model

nn.Module

Model to analyze

Returns: int

int

nemo_automodel.components.utils.model_utils.enable_radio_vit_fused_attn(
    model
)

Route RADIO ViT attention through F.scaled_dot_product_attention.

RADIO’s timm Attention blocks default to fused_attn=False, which materializes the full (B, H, seq, seq) attention tensor (~5 GiB per block at RADIO-v2-H + dynamic-resolution patch counts). Flipping fused_attn=True matches the Megatron-Bridge path which sets vision_config.use_flash_attn=True via attn_implementation="flash_attention_2".

No-op when the model has no RADIO vision tower.

Parameters:

model

The model to patch in place.

nemo_automodel.components.utils.model_utils.filter_forward_kwargs(
    model: torch.nn.Module,
    kwargs: dict
) -> dict

Drop kwargs that model.forward does not accept.

If the model exposes **kwargs or its signature cannot be inspected, the input kwargs are returned unchanged. The original dict is never mutated.

nemo_automodel.components.utils.model_utils.freeze_deepseek_v4_indexer_params(
    model
)

Freeze DeepSeek V4 indexer params that only feed discrete top-k masks.

nemo_automodel.components.utils.model_utils.freeze_minimax_m3_indexer_params(
    model
)

Freeze MiniMax M3 lightning-indexer params that only feed discrete top-k masks.

nemo_automodel.components.utils.model_utils.freeze_unused_kv_sharing_params(
    model
)

Freeze dead K/V parameters in KV-shared layers.

Models like Gemma4 E2B/E4B use KV-sharing where the last N layers reuse key/value states from earlier layers. The k_proj, v_proj, k_norm, and v_norm modules still exist in those shared layers but are never used during forward. Their parameters therefore receive no gradients, yet the optimizer still tracks them. On checkpoint resume the distributed checkpoint framework expects optimizer state for every parameter the optimizer was created with, but zero-gradient params may have been excluded from the saved state — causing a RuntimeError.

Calling this function before optimizer creation sets requires_grad=False on the dead parameters so the optimizer never tracks them, keeping save and load consistent.

Parameters:

model

The model (or pipeline-parallel model part).

nemo_automodel.components.utils.model_utils.get_lm_head_module(
    model: torch.nn.Module
) -> torch.nn.Module | None

Return the model’s LM head module, if one can be found.

nemo_automodel.components.utils.model_utils.get_lm_head_weight(
    model: torch.nn.Module
) -> torch.Tensor

Return the model’s LM-head weight, materializing DTensor weights when needed.

nemo_automodel.components.utils.model_utils.init_empty_weights()

A context manager under which models are initialized with all parameters on the specified device.

Example:

import torch.nn as nn
from nemo_automodel.components.utils.model_utils import init_empty_weights
with init_empty_weights():
    tst = nn.Linear(100, 100)  # on `cuda` device

Parameters:

device

`torch.device`

Device to initialize all parameters on.

nemo_automodel.components.utils.model_utils.parse_freeze_config(
    config: nemo_automodel.components.utils.model_utils.FreezeConfig | collections.abc.Mapping[str, typing.Any] | None
) -> nemo_automodel.components.utils.model_utils.FreezeConfig | None

Validate raw freeze_config YAML data and return a typed FreezeConfig.

Parameters:

config

FreezeConfig | Mapping[str, Any] | None

A FreezeConfig (returned unchanged), a raw mapping from the recipe config, or None.

Returns: FreezeConfig | None

The validated FreezeConfig, or None when no freeze configuration was

Raises:

TypeError: If config is neither a mapping nor a FreezeConfig.
ValueError: If config contains unknown options, malformed selectors, or non-boolean legacy options.

nemo_automodel.components.utils.model_utils.print_trainable_parameters(
    model: torch.nn.Module,
    name: str = 'Model'
) -> tuple[int, int]

Print the number of trainable parameters in the model.

Parameters:

model

nn.Module

Model to analyze

name

strDefaults to 'Model'

Label for the summary header (e.g. "Draft" to distinguish the draft model from the target in speculative-decoding training).

Returns: int

int

nemo_automodel.components.utils.model_utils.resolve_trust_remote_code(
    pretrained_model_name_or_path
)

Whitelist NVIDIA models to allow remote code execution.

Parameters:

pretrained_model_name_or_path

str

The name or path of the pretrained model.

Returns:

True if the model should be loaded with trust_remote_code, False otherwise.

nemo_automodel.components.utils.model_utils.skip_random_init()

Context manager to skip random weight initialization when loading pretrained models.

nemo_automodel.components.utils.model_utils.squeeze_input_for_thd(
    input_ids,
    position_ids,
    padding_mask,
    attn_kwargs,
    seqlens_padding_value = -1000
)

Squeeze batch dimension and prepare inputs for THD (total, hidden, depth) format.

This function removes the batch dimension from input tensors and processes attention kwargs for use with Transformer Engine’s THD format. It’s typically used when the batch has already been converted to THD format (with batch_size=1 as a placeholder dimension) and that dimension needs to be removed.

The function performs three key operations:

Removes the batch dimension (dim 0) from input tensors
Filters out padding values from cumulative sequence length tensors
Converts max_seqlen from tensor to scalar if needed

Parameters:

input_ids

torch.Tensor or None

Input token IDs with shape [1, total_tokens] or [1, total_tokens, hidden_dim]. The first dimension will be squeezed. None is permitted when the caller is feeding the model via inputs_embeds instead — embeddings are squeezed inside the model forward (the squeezed_for_thd branch in NemotronHModel.forward and analogous code paths), so this helper has nothing to squeeze and simply returns None for the input_ids slot.

position_ids

torch.Tensor

Position IDs with shape [1, total_tokens]. The first dimension will be squeezed.

padding_mask

torch.Tensor

Padding mask with shape [1, total_tokens]. The first dimension will be squeezed.

attn_kwargs

dict

Dictionary of attention-related tensors. May contain:

cu_seqlens: Cumulative sequence lengths [1, num_seqs+1]
cu_seqlens_padded: Cumulative padded sequence lengths [1, num_seqs+1]
max_seqlen: Maximum sequence length (tensor or int)
Other attention parameters (will be squeezed if tensors)

seqlens_padding_value

intDefaults to -1000

Sentinel value used to indicate padding in cu_seqlens and cu_seqlens_padded tensors. These values will be filtered out. Default: -1000.

Returns:

A tuple containing:

input_ids (torch.Tensor): Input IDs with batch dimension removed [total_tokens] or [total_tokens, hidden_dim]
position_ids (torch.Tensor): Position IDs with batch dimension removed [total_tokens]
padding_mask (torch.Tensor): Padding mask with batch dimension removed [total_tokens]
attn_kwargs (dict): Updated attention kwargs with:
- Batch dimensions removed from all tensor values
- Padding values filtered from cu_seqlens and cu_seqlens_padded
- max_seqlen converted to scalar if it was a tensor

nemo_automodel.components.utils.model_utils.VLM_INPUT_KEYS: tuple[str, ...] = ('input_ids', 'pixel_values', 'image_flags', 'imgs_sizes', 'image_position_ids',...

nemo_automodel.components.utils.model_utils._FREEZE_CONFIG_OPTIONS = {'freeze_modules', 'unfreeze_modules', 'freeze_vision_tower', 'freeze_audio_towe...

nemo_automodel.components.utils.model_utils.logger = logging.getLogger(__name__)