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

## Module Contents

### Classes

| Name                                                                                                | Description                                                                |
| --------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`CheckpointLoadPart`](#nemo_automodel-components-checkpoint-state_dict_adapter-CheckpointLoadPart) | One part of a checkpoint that is loaded and finished before the next part. |
| [`StateDictAdapter`](#nemo_automodel-components-checkpoint-state_dict_adapter-StateDictAdapter)     | Abstract base class for state dict transformations.                        |

### API

```python
class nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart(
    checkpoint_tensors: dict[str, torch.Tensor],
    model_keys: frozenset[str],
    temporary_checkpoint_keys: frozenset[str],
    finish: collections.abc.Callable[[], None]
)
```

Dataclass

One part of a checkpoint that is loaded and finished before the next part.

DCP fills `checkpoint_tensors` using their Hugging Face names and layouts. `finish` then converts any
temporary checkpoint tensors into the final model tensors named by `model_keys` before the loader advances to
the next part. Tensors that already have the model's exact dtype and layout may point directly at model storage.

**`checkpoint_tensors`** `dict[str, Tensor]`

Hugging Face checkpoint names mapped to DCP destinations. Destinations may use final model
storage or temporary storage owned by this part.

---

**`finish`** `Callable[[], None]`

Callback that installs temporary checkpoint tensors into final model storage. It must not save those
temporary tensors anywhere outside this load part.

---

**`model_keys`** `frozenset[str]`

Native model state-dict names completed after `finish` returns.

---

**`temporary_checkpoint_keys`** `frozenset[str]`

Names in `checkpoint_tensors` whose destinations use temporary storage rather
than final model storage. The loader uses these keys to report the largest per-rank temporary allocation.

---

```python
class nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter()
```

Abstract

Abstract base class for state dict transformations.

This class defines the interface for converting between native model
state dict format and other model state dict formats.

Most custom models need an adapter. Models whose HF weight names and tensor
layouts already match can omit it.

**`_supports_low_memory_dcp_load`** `bool = False`

---

**`checkpoint_key_aliases`** `dict[str, str]`

Current HF names mapped to legacy names accepted during strict checkpoint reads.

The loader uses an alias only when the current name is absent, then restores the current name
before invoking from\_hf. Saving always uses the current names.

---

**`supports_low_memory_dcp_load`** `bool`

Whether DCP can load the checkpoint with zero or small temporary tensors.

Most checkpoint tensors must load directly into the model's existing weight memory. Small temporary tensors
are allowed when they are converted and released after the read. Enable this only when the extra memory is
safely below a full model copy, including when loading on one GPU.

---

```python
nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: typing.Any,
    kwargs = {}
) -> list[tuple[str, typing.Any]]
```

abstract

Convert a single tensor from native format to HuggingFace format.

**Parameters:**

**`fqn`** `str`

Fully qualified name of the tensor in native format

---

**`tensor`** `Any`

The tensor to convert

---

**`**kwargs`** — default: \{}

Additional arguments for conversion

---

**Returns:** `list[tuple[str, Any]]`

List of (fqn, tensor) tuples in HuggingFace format.

```python
nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
    kwargs = {}
) -> dict[str, typing.Any]
```

abstract

Obtain native model state dict from HuggingFace format.

**Parameters:**

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

The HuggingFace format state dict

---

**`device_mesh`** `Optional[DeviceMesh]` — default: None

Optional device mesh for DTensor expert parallelism.
If provided, only loads experts needed for the current rank.

---

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

The converted native model state dict

```python
nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter.get_hf_state_dict_keys(
    state_dict: dict[str, typing.Any]
) -> list[str]
```

Return the Hugging Face keys produced by `to_hf` without converting real weights.

**Parameters:**

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

Native model state mapping. Tensor values may have
arbitrary rank and axis order and retain their exact parameter
or buffer layouts.

---

**Returns:** `list[str]`

Hugging Face state-dict keys in adapter iteration order.

```python
nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter.iter_checkpoint_load_parts(
    model_state_dict: dict[str, torch.Tensor],
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart] | None
```

Optionally load and finish a quantized checkpoint in small parts.

Ordinary adapters do not need to implement this method. It is only for adapters that cannot let DCP write
checkpoint tensors directly into model storage because a dtype or layout conversion is required, but can
complete that conversion for one small part of the model at a time.

**Parameters:**

**`model_state_dict`** `dict[str, torch.Tensor]`

Native model names mapped to the final parameter and persistent-buffer tensors that the
checkpoint must populate. Tensors retain their model-specific shapes, axis orders, dtypes, devices,
strides, distributed placements, and storage.

---

**`device_mesh`** `Optional[DeviceMesh]` — default: None

Optional device mesh describing distributed model storage.

---

**Returns:** `Iterator[CheckpointLoadPart] | None`

An iterator of dependency-complete load parts, or `None` when the adapter does not support this path.

```python
nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter.map_peft_target_module_to_hf(
    name: str,
    v4_compatible: bool = False
) -> str
```

Translate a PEFT target-module name to the HuggingFace layout.

adapter\_config.json's target\_modules are collected from native module
names. Adapters whose `to_hf` renames modules (e.g. Kimi K3's
`mlp.experts.{E}.gate_proj` -> `block_sparse_moe.experts.{E}.w1`)
should override this with the same renames so PEFT can resolve the
entries against the converted checkpoint.

**Parameters:**

**`name`** `str`

A target-module name in native layout.

---

**`v4_compatible`** `bool` — default: False

Whether to target the legacy Transformers v4 layout.

---

**Returns:** `str`

The name in HuggingFace layout. Defaults to the name unchanged.

```python
nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    kwargs = {}
) -> dict[str, typing.Any]
```

abstract

Convert from native model state dict to HuggingFace format.

**Parameters:**

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

The native model state dict

---

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

The converted HuggingFace format state dict