nemo_automodel.components.models.step5.state_dict_adapter

Step-5 FP8 loading, interleaved QKV mapping and native expert conversion.

Module Contents

Classes

Name	Description
`Step5StateDictAdapter`	Preserve Step-family expert conversion and add Step-5 attention mappings.

Data

_DECODER_LAYER

_INDEXER_NAMES

API

class nemo_automodel.components.models.step5.state_dict_adapter.Step5StateDictAdapter(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)

Bases: Step3p5StateDictAdapter

Preserve Step-family expert conversion and add Step-5 attention mappings.

Loading needs bounded per-projection conversion storage, so the adapter does not advertise the DCP direct-destination fast path.

_load_fp8_experts

= quantization is not None

_unquantized_expert_modules

= frozenset(excluded)

nemo_automodel.components.models.step5.state_dict_adapter.Step5StateDictAdapter._iter_checkpoint_load_parts(
    model_state_dict: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart]

Load mixed FP8/BF16 expert projections into final BF16 storage.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Final native storage: fused gate/up tensors [experts, hidden, 2 * intermediate], down tensors [experts, intermediate, hidden], and other native model tensors. DTensors shard experts over ep and input dimensions over ep_shard.

device_mesh

DeviceMesh | NoneDefaults to None

Optional expert mesh used for native/HF conversion.

nemo_automodel.components.models.step5.state_dict_adapter.Step5StateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue = {}
) -> list[tuple[str, torch.Tensor]]

Export native attention and expert tensors to checkpoint names.

Parameters:

fqn

str

Native tensor name.

tensor

torch.Tensor

Tensor with the native layout documented by from_hf.

exclude_key_regex

str | NoneDefaults to None

Optional exported-name exclusion regex.

quantization

boolDefaults to False

Whether quantized export is requested (unsupported).

**kwargs

ConfigValueDefaults to {}

Checkpointer options forwarded to the base adapter.

Returns: list[tuple[str, torch.Tensor]]

Named source tensors with [out, in] projections or [experts, out, in] experts.

nemo_automodel.components.models.step5.state_dict_adapter.Step5StateDictAdapter.from_hf(
    hf_state_dict: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue = {}
) -> dict[str, torch.Tensor]

Load source tensors into trainable, interleaved native storage.

Parameters:

hf_state_dict

dict[str, torch.Tensor]

Mapping to source tensors. Experts have layout [experts, out, in]; Q/K/V have layout [heads * head_dim, hidden].

device_mesh

DeviceMesh | NoneDefaults to None

Optional expert mesh handled by the Step-family adapter.

**kwargs

ConfigValueDefaults to {}

Checkpointer options forwarded to the base adapter.

Returns: dict[str, torch.Tensor]

Mapping to native tensors. QKV has layout

nemo_automodel.components.models.step5.state_dict_adapter.Step5StateDictAdapter.iter_checkpoint_load_parts(
    model_state_dict: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart] | None

Load in bounded parts unless this rank holds only part of the decoder.

Every rank issues one load collective per part. Pipeline stages own different decoder layers, so their part counts differ; such ranks use the ordinary checkpoint load instead.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Final native storage; see _iter_checkpoint_load_parts.

device_mesh

DeviceMesh | NoneDefaults to None

Optional expert mesh used for native/HF conversion.

Returns: Iterator[CheckpointLoadPart] | None

Checkpoint load parts, or None on a distributed rank whose state

nemo_automodel.components.models.step5.state_dict_adapter._DECODER_LAYER = re.compile('^model\\.layers\\.(\\d+)\\.')

nemo_automodel.components.models.step5.state_dict_adapter._INDEXER_NAMES = {'sparse_indexer_q': 'indexer.q_proj', 'sparse_indexer_k': 'indexer.k_proj', 'sp...