> 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.                            |
| [`_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.                                         |
| [`_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.                          |

### Data

[`_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
)
```

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:**

The PyTorch model to track.

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

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.

Whether missing or unexpected keys should fail the load.

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's state dict with CPU offloading enabled.

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

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:**

The PyTorch model associated with the optimizer.

The optimizer to track.

Optional learning rate 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 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 with CPU offloading enabled.

```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._get_lm_head_weight_and_name(
    model: torch.nn.Module
) -> typing.Optional[tuple[torch.Tensor, str]]
```

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