> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.utils.model_utils

## Module Contents

### Classes

| Name                                                                            | Description                                                                |
| ------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`FreezeConfig`](#nemo_automodel-components-utils-model_utils-FreezeConfig)     | Typed schema for the `freeze_config` recipe section.                       |
| [`ModuleSelector`](#nemo_automodel-components-utils-model_utils-ModuleSelector) | Select modules by exact canonical path or case-sensitive shell-style glob. |

### Functions

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

### Data

[`VLM_INPUT_KEYS`](#nemo_automodel-components-utils-model_utils-VLM_INPUT_KEYS)

[`_FREEZE_CONFIG_OPTIONS`](#nemo_automodel-components-utils-model_utils-_FREEZE_CONFIG_OPTIONS)

[`logger`](#nemo_automodel-components-utils-model_utils-logger)

### API

```python
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:

1. Full fine-tuning preserves the model's existing trainability state;
   PEFT establishes a LoRA-trainable, base-frozen baseline.
2. `freeze_modules` recursively freezes the selected modules.
3. `unfreeze_modules` recursively unfreezes the selected modules and wins
   on overlap.
4. Framework-required freezes (dead K/V projections, indexer parameters)
   are applied by the infrastructure afterwards and remain protected.
5. 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`

---

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

Validate selectors and legacy compatibility options.

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

Return whether either generic selector field was explicitly declared.

```python
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
`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`

---

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

Validate that exactly one matching mode is configured.

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

Return the selector in its `key: value` configuration form.

```python
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`.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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

```python
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.

```python
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.

```python
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).

```python
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.

**Legacy modality booleans**

* freeze\_vision\_tower: bool (default True in legacy-only configurations)
* freeze\_audio\_tower: bool (default False)
* freeze\_language\_model: bool (default False)
* freeze\_video\_embedder: bool (default False)

**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`** `bool` — default: 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.

---

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

Cast fp32 parameters and buffers to bf16 for FSDP2 compatibility.

```python
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

```python
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.

---

```python
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.

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

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

```python
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.

```python
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).

---

```python
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.

```python
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.

```python
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:

```python
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.

---

```python
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.

```python
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`** `str` — default: 'Model'

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

---

**Returns:** `int`

int

```python
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.

```python
nemo_automodel.components.utils.model_utils.skip_random_init()
```

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

```python
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 four key operations:

1. Removes the batch dimension (dim 0) from input tensors
2. Filters out padding values from cumulative sequence length tensors
3. Converts max\_seqlen from tensor to scalar if needed
4. Resolves the microbatch's TE padding flag from its complete boundary arrays

**Example**

\>>> input\_ids = torch.tensor(\[\[1, 2, 3, 4, 5]])  # \[1, 5]
\>>> position\_ids = torch.tensor(\[\[0, 1, 2, 3, 4]])  # \[1, 5]
\>>> padding\_mask = torch.tensor(\[\[False, False, False, False, False]])  # \[1, 5]
\>>> attn\_kwargs = \{
...     'cu\_seqlens': torch.tensor(\[\[0, 3, 5, -1000]]),  # \[1, 4] with padding
...     'cu\_seqlens\_padded': torch.tensor(\[\[0, 3, 5, -1000]]),
...     'max\_seqlen': torch.tensor(\[3])
... }
\>>> ids, pos, mask, kwargs = squeeze\_input\_for\_thd(
...     input\_ids, position\_ids, padding\_mask, attn\_kwargs
... )
\>>> ids.shape
torch.Size(\[5])
\>>> kwargs\['cu\_seqlens']  # Padding value filtered out
tensor(\[0, 3, 5])
\>>> kwargs\['max\_seqlen']  # Converted to scalar
3

**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`** `int` — default: -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

> **Note**
>
> This function modifies attn\_kwargs in-place. If you need to preserve the original
> dictionary, pass a copy.

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

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

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