nemo_automodel.components.models.inkling.state_dict_adapter

State-dict adapter for the native Thinking Machines Inkling checkpoint.

Module Contents

Classes

Name	Description
`InklingStateDictAdapter`	Convert Inkling weights between native AutoModel and raw checkpoint layouts.

Functions

Name	Description
`_interleave`	Match Transformers’ Inkling `Interleave` checkpoint conversion.
`_native_to_raw_key`	-
`_raw_to_native_key`	-

Data

_RAW_TRANSPOSED_SUFFIXES

API

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

Bases: StateDictAdapter

Convert Inkling weights between native AutoModel and raw checkpoint layouts.

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

Convert one native tensor to the raw Inkling checkpoint layout.

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

Load raw Inkling keys or an already-converted Transformers state dict.

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

Expose raw Inkling checkpoint keys, using views for large matrices.

nemo_automodel.components.models.inkling.state_dict_adapter._interleave(
    tensor: torch.Tensor,
    dim: int,
    inverse: bool = False
) -> torch.Tensor

Match Transformers’ Inkling Interleave checkpoint conversion.

nemo_automodel.components.models.inkling.state_dict_adapter._native_to_raw_key(
    key: str
) -> str

nemo_automodel.components.models.inkling.state_dict_adapter._raw_to_native_key(
    key: str
) -> str

nemo_automodel.components.models.inkling.state_dict_adapter._HF_DENSE_RE = re.compile('((?:model\\.)?language_model\\.layers\\.\\d+\\.mlp)\\.(gate_proj|up_...

nemo_automodel.components.models.inkling.state_dict_adapter._HF_EXPERT_RE = re.compile('(?:model\\.)?language_model\\.layers\\.(\\d+)\\.mlp\\.experts\\.(gat...

nemo_automodel.components.models.inkling.state_dict_adapter._HF_SHARED_RE = re.compile('((?:model\\.)?language_model\\.layers\\.\\d+\\.mlp\\.shared_experts)...

nemo_automodel.components.models.inkling.state_dict_adapter._RAW_EXPERT_RE = re.compile('model\\.llm\\.layers\\.(\\d+)\\.mlp\\.experts\\.(w13_weight|w2_weigh...

nemo_automodel.components.models.inkling.state_dict_adapter._RAW_TRANSPOSED_SUFFIXES = ('.mlp.experts.w13_weight', '.mlp.experts.w2_weight', '.mlp.shared_experts.share...