nemo_automodel.components.models.kimi_k3.state_dict_adapter

State-dict conversion for native Kimi K3 and its MXFP4 checkpoint.

Module Contents

Classes

Name	Description
`KimiK3StateDictAdapter`	Convert K3 split/packed experts to AutoModel grouped experts.

Functions

Name	Description
`_hf_moe_key_to_native`	Rewrite a checkpoint MoE key onto the decoder block’s `mlp` submodule.
`_native_moe_key_to_hf`	Inverse of :func:`_hf_moe_key_to_native`.
`_route_kda_fp32_holder`	-
`_strip_kda_fp32_holder`	-
`_upcast_fp32_state_tensor`	-
`dequantize_mxfp4`	Decode K3 `[out, in / 2]` MXFP4 bytes into `[out, in]` weights.

Data

_FP32_KEY_PARTS

_GENERIC_TO_HF_EXPERT_PROJ

_HF_TO_GENERIC_EXPERT_PROJ

_KDA_FP32_HOLDER

_KDA_FP32_OP_HOLDER

_KDA_FP32_OP_PARAM

_KDA_FP32_PARAM_NAMES

_MOE_CHILD_SEGMENTS

_MXFP4_VALUES

API

class nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.float32
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert K3 split/packed experts to AutoModel grouped experts.

HF stores routed experts as per-expert Kimi names:

block_sparse_moe.experts.{E}.w1.weight: SwiGLU gate projection, shape [inter, hidden].
block_sparse_moe.experts.{E}.w3.weight: SwiGLU up projection, shape [inter, hidden].
block_sparse_moe.experts.{E}.w2.weight: down projection, shape [hidden, inter].

Automodel stores grouped experts as:

mlp.experts.gate_and_up_projs with shape [experts, hidden, 2 * inter].
mlp.experts.down_projs with shape [experts, inter, hidden].

The decoder block names both its dense and its MoE feed-forward mlp (the naming the custom-MoE parallelizer looks for), so the checkpoint’s block_sparse_moe path segment is translated here in both directions.

_expert_path_segment

str

config

= getattr(config, 'text_config', config)

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._add_hf_text_prefix(
    key: str
) -> str

staticmethod

Prefix native text keys while leaving vision/projector keys unchanged.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._dequantize_packed_experts(
    state_dict: dict[str, typing.Any]
) -> None

Decode packed experts, writing directly into model views when available.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._make_mxfp4_load_destinations(
    key: str,
    weight: torch.Tensor
) -> list[tuple[str, torch.Tensor]]

staticmethod

Create packed checkpoint destinations matching one plain [out, in] expert weight.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._map_generic_expert_key_to_hf(
    key: str
) -> str

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._map_hf_expert_key_to_generic(
    key: str
) -> str

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._normalize_checkpoint_tensor(
    key: str,
    value: typing.Any
) -> typing.Any

Remove K3’s zero padding from the per-head KDA decay parameter.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._pad_checkpoint_a_log(
    key: str,
    value: typing.Any
) -> typing.Any

Restore the 128-entry checkpoint storage layout for KDA A_log.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._split_experts_weights(
    weight: torch.Tensor,
    n_experts: int
) -> list[torch.Tensor]

Split grouped experts, tolerating DTensors whose mesh dim is not named ep.

Parameters:

weight

torch.Tensor

Grouped routed-expert tensor of shape [experts, …]. May be a plain tensor or a DTensor sharded or replicated over the expert axis; for Shard(0), the local shard covers this rank’s expert slice.

n_experts

int

Global number of routed experts.

Returns: list[torch.Tensor]

List of per-expert tensors of shape […] for the experts local to this rank.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter._strip_hf_text_prefix(
    key: str
) -> str

staticmethod

Remove the K3 checkpoint’s language_model. namespace.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: typing.Any,
    kwargs: typing.Any = {}
) -> list[tuple[str, typing.Any]]

Convert one Automodel tensor to one or more Kimi HF tensors.

Parameters:

fqn

str

Fully qualified native tensor name.

tensor

Any

Native tensor. Grouped routed expert tensors use [experts, hidden, 2 * inter] for gate/up and [experts, inter, hidden] for down.

**kwargs

AnyDefaults to {}

Adapter options forwarded by checkpoint save/load.

Returns: list[tuple[str, Any]]

HF key/tensor pairs. Split expert tensors use Kimi w1/w2/w3 names.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    kwargs: typing.Any = {}
) -> dict[str, typing.Any]

Convert Kimi HF checkpoint keys to Automodel native keys.

Parameters:

hf_state_dict

dict[str, Any]

HF state dict whose routed expert tensors use split Kimi names.

device_mesh

DeviceMesh | NoneDefaults to None

Optional EP/FSDP mesh used to load only local expert shards.

**kwargs

AnyDefaults to {}

Adapter options forwarded by checkpoint load.

Returns: dict[str, Any]

Native state dict with grouped routed expert tensors.

nemo_automodel.components.models.kimi_k3.state_dict_adapter.KimiK3StateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    kwargs: typing.Any = {}
) -> dict[str, typing.Any]

Convert Automodel native tensors to Kimi HF checkpoint keys.

nemo_automodel.components.models.kimi_k3.state_dict_adapter._hf_moe_key_to_native(
    key: str
) -> str

Rewrite a checkpoint MoE key onto the decoder block’s mlp submodule.

nemo_automodel.components.models.kimi_k3.state_dict_adapter._native_moe_key_to_hf(
    key: str
) -> str

Inverse of :func:_hf_moe_key_to_native.

nemo_automodel.components.models.kimi_k3.state_dict_adapter._route_kda_fp32_holder(
    key: str
) -> str

nemo_automodel.components.models.kimi_k3.state_dict_adapter._strip_kda_fp32_holder(
    key: str
) -> str

nemo_automodel.components.models.kimi_k3.state_dict_adapter._upcast_fp32_state_tensor(
    key: str,
    value: typing.Any
) -> typing.Any

nemo_automodel.components.models.kimi_k3.state_dict_adapter.dequantize_mxfp4(
    weight_packed: torch.Tensor,
    weight_scale: torch.Tensor,
    dtype: torch.dtype = torch.bfloat16
) -> torch.Tensor

Decode K3 [out, in / 2] MXFP4 bytes into [out, in] weights.

nemo_automodel.components.models.kimi_k3.state_dict_adapter._FP32_KEY_PARTS = ('A_log', 'dt_bias', 'e_score_correction_bias', 'q_conv1d.weight', 'k_conv1d.wei...

nemo_automodel.components.models.kimi_k3.state_dict_adapter._GENERIC_TO_HF_EXPERT_PROJ = {value: key for key, value in (_HF_TO_GENERIC_EXPERT_PROJ.items())}

nemo_automodel.components.models.kimi_k3.state_dict_adapter._HF_TO_GENERIC_EXPERT_PROJ = {'w1': 'gate_proj', 'w2': 'down_proj', 'w3': 'up_proj'}

nemo_automodel.components.models.kimi_k3.state_dict_adapter._KDA_FP32_HOLDER = re.compile('(\\.self_attn)\\._fp32_params\\.')

nemo_automodel.components.models.kimi_k3.state_dict_adapter._KDA_FP32_OP_HOLDER = re.compile('(\\.self_attn\\.(?:q_conv1d|k_conv1d|v_conv1d|o_norm))\\._fp32_param...

nemo_automodel.components.models.kimi_k3.state_dict_adapter._KDA_FP32_OP_PARAM = re.compile('(\\.self_attn\\.(?:q_conv1d|k_conv1d|v_conv1d|o_norm))\\.(weight)$')

nemo_automodel.components.models.kimi_k3.state_dict_adapter._KDA_FP32_PARAM_NAMES = ('A_log', 'dt_bias')

nemo_automodel.components.models.kimi_k3.state_dict_adapter._MOE_CHILD_SEGMENTS = ('gate.', 'shared_experts.', 'routed_expert_down_proj.', 'routed_expert_up_proj....

nemo_automodel.components.models.kimi_k3.state_dict_adapter._MXFP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4....