nemo_automodel.components.checkpoint.state_dict_adapter

Module Contents

Classes

Name	Description
`StateDictAdapter`	Abstract base class for state dict transformations.

API

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.

_supports_write_through_checkpoint_load

bool = False

supports_write_through_checkpoint_load

bool

Whether all checkpoint-load destinations write through to final model storage.

Adapters should set _supports_write_through_checkpoint_load only when every tensor produced by to_hf for base-checkpoint loading is either the original model tensor or a view that writes through to it. The checkpoint loader uses this guarantee to avoid materializing a converted full state dict on the host.

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.

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