nemo_automodel.components.checkpoint.state_dict_adapter

Module Contents

Classes

Name	Description
`CheckpointLoadPart`	One part of a checkpoint that is loaded and finished before the next part.
`StateDictAdapter`	Abstract base class for state dict transformations.

API

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]

finish

Callable[[], None]

model_keys

frozenset[str]

temporary_checkpoint_keys

frozenset[str]

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

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.

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

Defaults to {}

Additional arguments for conversion

Returns: list[tuple[str, Any]]

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

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]Defaults to 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

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.

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]Defaults to 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.

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

boolDefaults to False

Whether to target the legacy Transformers v4 layout.

Returns: str

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

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