> 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.checkpoint.stateful_wrappers

## Module Contents

### Classes

| Name                                                                                       | Description                                                             |
| ------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------- |
| [`ModelState`](#nemo_automodel-components-checkpoint-stateful_wrappers-ModelState)         | Helper class for tracking model state in distributed checkpointing.     |
| [`OptimizerState`](#nemo_automodel-components-checkpoint-stateful_wrappers-OptimizerState) | Helper class for tracking optimizer state in distributed checkpointing. |

### Functions

| Name                                                                                                                             | Description                                                                                 |
| -------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- |
| [`_add_outer_prefix`](#nemo_automodel-components-checkpoint-stateful_wrappers-_add_outer_prefix)                                 | Prepend `prefix` once to every key in-place (inverse of `_drop_outer_prefix`).              |
| [`_drop_outer_prefix`](#nemo_automodel-components-checkpoint-stateful_wrappers-_drop_outer_prefix)                               | Remove the *first* occurrence of `prefix` on every key in-place.                            |
| [`_gather_peft_state_dict_across_pp`](#nemo_automodel-components-checkpoint-stateful_wrappers-_gather_peft_state_dict_across_pp) | All-gather PEFT adapter tensors across a pipeline-parallel group.                           |
| [`_get_lm_head_weight_and_name`](#nemo_automodel-components-checkpoint-stateful_wrappers-_get_lm_head_weight_and_name)           | -                                                                                           |
| [`_get_peft_state_dict`](#nemo_automodel-components-checkpoint-stateful_wrappers-_get_peft_state_dict)                           | Extract only trainable PEFT adapter weights, bypassing DCP.                                 |
| [`_has_expert_parallelism`](#nemo_automodel-components-checkpoint-stateful_wrappers-_has_expert_parallelism)                     | Check if any MoE expert module in the model has expert parallelism enabled.                 |
| [`_has_quantized_params`](#nemo_automodel-components-checkpoint-stateful_wrappers-_has_quantized_params)                         | Check if model has any BitsAndBytes quantized modules.                                      |
| [`_is_quantized_module`](#nemo_automodel-components-checkpoint-stateful_wrappers-_is_quantized_module)                           | Check if a module is a BitsAndBytes quantized type.                                         |
| [`_materialize_missing_adam_state`](#nemo_automodel-components-checkpoint-stateful_wrappers-_materialize_missing_adam_state)     | Create zero-valued Adam state for parameters that do not have state yet.                    |
| [`_rename_dora_keys_from_hf`](#nemo_automodel-components-checkpoint-stateful_wrappers-_rename_dora_keys_from_hf)                 | Reverse of \_rename\_dora\_keys\_to\_hf: convert HF PEFT key format back to internal names. |
| [`_rename_dora_keys_to_hf`](#nemo_automodel-components-checkpoint-stateful_wrappers-_rename_dora_keys_to_hf)                     | Rename DoRA magnitude keys to match HF PEFT's saved checkpoint format in-place.             |
| [`_safe_op_set_extra_state`](#nemo_automodel-components-checkpoint-stateful_wrappers-_safe_op_set_extra_state)                   | -                                                                                           |
| [`_safe_set_extra_state`](#nemo_automodel-components-checkpoint-stateful_wrappers-_safe_set_extra_state)                         | -                                                                                           |
| [`_set_peft_state_dict`](#nemo_automodel-components-checkpoint-stateful_wrappers-_set_peft_state_dict)                           | Load trainable PEFT adapter weights into the model, bypassing DCP.                          |
| [`_zeros_like_optimizer_param`](#nemo_automodel-components-checkpoint-stateful_wrappers-_zeros_like_optimizer_param)             | Allocate zero optimizer state matching a parameter.                                         |

### Data

[`_OPTIMIZER_PARTS_KEY`](#nemo_automodel-components-checkpoint-stateful_wrappers-_OPTIMIZER_PARTS_KEY)

[`_OPTIMIZER_PART_KEY_PREFIX`](#nemo_automodel-components-checkpoint-stateful_wrappers-_OPTIMIZER_PART_KEY_PREFIX)

[`_PREFIX`](#nemo_automodel-components-checkpoint-stateful_wrappers-_PREFIX)

[`_original_op_set_extra_state`](#nemo_automodel-components-checkpoint-stateful_wrappers-_original_op_set_extra_state)

[`_original_set_extra_state`](#nemo_automodel-components-checkpoint-stateful_wrappers-_original_set_extra_state)

### API

```python
class nemo_automodel.components.checkpoint.stateful_wrappers.ModelState(
    model: torch.nn.Module | list[torch.nn.Module],
    is_peft: bool = False,
    is_init_step: bool = False,
    skip_task_head_prefixes: list[str] | None = None,
    cpu_offload: bool = False,
    pp_group: torch.distributed.ProcessGroup | None = None,
    has_expert_parallelism: bool = False
)
```

Helper class for tracking model state in distributed checkpointing.

This class is compliant with the Stateful protocol, allowing DCP to automatically
call state\_dict/load\_state\_dict as needed in the dcp.save/load APIs.

**Parameters:**

**`model`** `torch.nn.Module | list[torch.nn.Module]`

The PyTorch model to track.

---

**`has_local_tied_lm_head`** `= has_local_tied_lm_head(self.model[0])`

---

**`model`**

---

**`skip_task_head_prefixes`** `= skip_task_head_prefixes or []`

---

**`uses_tied_lm_head`** `= is_tied_word_embeddings(self.model[0])`

---

```python
nemo_automodel.components.checkpoint.stateful_wrappers.ModelState._get_base_model_state_dict() -> dict[str, typing.Any]
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers.ModelState._refresh_local_tied_lm_head() -> None
```

Refresh tied-head metadata after DCP has normalized module state.

```python
nemo_automodel.components.checkpoint.stateful_wrappers.ModelState._set_base_model_state_dict(
    state_dict: dict[str, typing.Any]
) -> None
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers.ModelState.load_state_dict(
    state_dict: dict[str, typing.Any],
    strict: bool = True,
    broadcast_from_rank0: bool = True
) -> None
```

Load the state dictionary into the model.

**Parameters:**

**`state_dict`** `dict[str, Any]`

Model state mapping whose tensor values may have arbitrary
rank and axis order and retain each parameter or buffer's exact
shape and DTensor placement.

---

**`strict`** `bool` — default: True

Whether missing or unexpected keys should fail the load.

---

**`broadcast_from_rank0`** `bool` — default: True

Whether rank 0 owns the full PEFT state dict.
Set to `False` when every rank in a model-local process group
loaded the adapter independently.

---

```python
nemo_automodel.components.checkpoint.stateful_wrappers.ModelState.state_dict() -> dict[str, typing.Any]
```

Get the model's state dictionary.

**Returns:** `dict[str, Any]`

Dictionary containing the model state dict, optionally offloaded to CPU.

```python
class nemo_automodel.components.checkpoint.stateful_wrappers.OptimizerState(
    model: torch.nn.Module | list[torch.nn.Module],
    optimizer: torch.optim.Optimizer | list[torch.optim.Optimizer],
    scheduler: typing.Any | None = None,
    is_peft: bool = False,
    cpu_offload: bool = False,
    has_expert_parallelism: bool = False,
    optimizer_part_ids: list[int] | None = None
)
```

Helper class for tracking optimizer state in distributed checkpointing.

This class is compliant with the Stateful protocol, allowing DCP to automatically
call state\_dict/load\_state\_dict as needed in the dcp.save/load APIs.

**Parameters:**

**`model`** `torch.nn.Module | list[torch.nn.Module]`

The PyTorch model associated with the optimizer.

---

**`optimizer`** `torch.optim.Optimizer | list[torch.optim.Optimizer]`

The optimizer to track.

---

**`scheduler`** `Any | None` — default: None

Optional learning rate scheduler.

---

**`_use_native_optimizer_state`**

---

**`model`**

---

**`optimizer`**

---

**`scheduler`**

---

```python
nemo_automodel.components.checkpoint.stateful_wrappers.OptimizerState.load_state_dict(
    state_dict: dict[str, typing.Any]
) -> None
```

Load the state dictionaries into the optimizer and scheduler.

**Parameters:**

**`state_dict`** `dict`

State dictionary containing optimizer and scheduler states to load.

---

```python
nemo_automodel.components.checkpoint.stateful_wrappers.OptimizerState.state_dict() -> dict[str, typing.Any]
```

Get the optimizer and scheduler state dictionaries.

**Returns:** `dict[str, Any]`

Dictionary containing the optimizer and scheduler state dicts, optionally offloaded to CPU.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._add_outer_prefix(
    sd: dict[str, typing.Any],
    prefix: str = _PREFIX,
    skip_keys: list[str] | None = None
) -> None
```

Prepend `prefix` once to every key in-place (inverse of `_drop_outer_prefix`).

```python
nemo_automodel.components.checkpoint.stateful_wrappers._drop_outer_prefix(
    sd: dict[str, typing.Any],
    prefix: str = _PREFIX
) -> None
```

Remove the *first* occurrence of `prefix` on every key in-place.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._gather_peft_state_dict_across_pp(
    local_state_dict: dict[str, typing.Any],
    pp_group: torch.distributed.ProcessGroup
) -> dict[str, typing.Any]
```

All-gather PEFT adapter tensors across a pipeline-parallel group.

Pipeline parallelism partitions the model's layers across PP ranks: each
rank's local module only contains its own stage's layers. The local
collection in :meth:`ModelState.state_dict` gathers PEFT tensors solely from
the local model parts, so under `pp_size > 1` the per-rank state dict is
missing every layer owned by another stage. Saving that directly yields a
truncated adapter (only \~1/pp of the layers), which silently degrades a
merged model.

This gathers the per-rank PEFT dicts over `pp_group` and merges them by FQN
so every rank returns the complete adapter. Keys are globally unique across PP
stages (layer indices never overlap between stages), so the union is exact and
order-independent; on the rare chance the same key appears on two ranks (e.g.
a replicated tied parameter) the lowest-rank value wins deterministically.

**Parameters:**

**`local_state_dict`** `dict[str, Any]`

This rank's PEFT tensors (already CPU, bf16/fp32).

---

**`pp_group`** `torch.distributed.ProcessGroup`

The pipeline-parallel process group to gather over.

---

**Returns:** `dict[str, Any]`

The merged PEFT state dict containing every PP stage's adapter tensors.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._get_lm_head_weight_and_name(
    model: torch.nn.Module
) -> tuple[torch.Tensor, str] | None
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers._get_peft_state_dict(
    model: torch.nn.Module
) -> dict[str, typing.Any]
```

Extract only trainable PEFT adapter weights, bypassing DCP.

This function directly iterates over model parameters to collect trainable weights,
avoiding PyTorch DCP's state\_dict traversal which fails on (1) BitsAndBytes quantized
modules (Params4bit, Int8Params, etc.) and (2) MoE models with expert parallelism
where expert weights are sharded across EP ranks.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._has_expert_parallelism(
    model: torch.nn.Module
) -> bool
```

Check if any MoE expert module in the model has expert parallelism enabled.

After EP initialization, expert modules (GroupedExpertsDeepEP, GroupedExpertsTE)
store `ep_size` on themselves. A value > 1 signals that expert weights are
sharded across EP ranks and DCP's state\_dict APIs cannot handle them.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._has_quantized_params(
    model: torch.nn.Module
) -> bool
```

Check if model has any BitsAndBytes quantized modules.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._is_quantized_module(
    module: torch.nn.Module
) -> bool
```

Check if a module is a BitsAndBytes quantized type.

Detects quantization by checking for `quant_state` attribute which is
common across BitsAndBytes quantized module types (Params4bit, Int8Params, etc.).

```python
nemo_automodel.components.checkpoint.stateful_wrappers._materialize_missing_adam_state(
    optimizer: torch.optim.Optimizer
) -> None
```

Create zero-valued Adam state for parameters that do not have state yet.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._rename_dora_keys_from_hf(
    sd: dict[str, typing.Any]
) -> None
```

Reverse of \_rename\_dora\_keys\_to\_hf: convert HF PEFT key format back to internal names.

Handles both the current on-disk format (`<module>.lora_magnitude_vector`)
and the legacy format that included `.default.weight` for robustness.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._rename_dora_keys_to_hf(
    sd: dict[str, typing.Any]
) -> None
```

Rename DoRA magnitude keys to match HF PEFT's saved checkpoint format in-place.

HF PEFT's `get_peft_model_state_dict` strips the adapter name and the
`.weight` suffix from `lora_magnitude_vector.<adapter>.<weight>` so the
round-trip format on disk is simply `<module>.lora_magnitude_vector`.
When loading, `set_peft_model_state_dict` re-inserts the adapter name
and the `.weight` suffix automatically, so we must NOT include them here.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._safe_op_set_extra_state(
    self,
    state
)
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers._safe_set_extra_state(
    self,
    state
)
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers._set_peft_state_dict(
    model: torch.nn.Module,
    state_dict: dict[str, typing.Any]
) -> None
```

Load trainable PEFT adapter weights into the model, bypassing DCP.

Mirrors \_get\_peft\_state\_dict: directly assigns saved tensors to model parameters
by name, handling DTensor re-sharding for EP-parallel weights. This avoids
DCP's set\_model\_state\_dict() which raises KeyError on expert-parallel FQNs.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._zeros_like_optimizer_param(
    param: torch.Tensor
) -> torch.Tensor
```

Allocate zero optimizer state matching a parameter.

**Parameters:**

**`param`** `torch.Tensor`

Tensor of arbitrary shape representing one optimizer parameter.

---

**Returns:** `torch.Tensor`

Zero tensor with the same shape, dtype, device, and layout as `param`.

```python
nemo_automodel.components.checkpoint.stateful_wrappers._OPTIMIZER_PARTS_KEY = 'optimizer_parts'
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers._OPTIMIZER_PART_KEY_PREFIX = 'stage_'
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers._PREFIX = 'model.'
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers._original_op_set_extra_state = te_ops.BasicOperation.set_extra_state
```

```python
nemo_automodel.components.checkpoint.stateful_wrappers._original_set_extra_state = te_base.TransformerEngineBaseModule.set_extra_state
```