Skip to navigation
DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsInkling2nemo_automodel.components.models.inkling2.state_dict_adapter

nemo_automodel.components.models.inkling2.state_dict_adapter

View as MarkdownOpen in Claude

Inkling2 text checkpoint mapping and mixed-precision import through DCP load targets.

Module Contents

Classes

NameDescription
Inkling2StateDictAdapterMap the delivered text weights; preserve floating checkpoint layouts on export.
_FloatingLayoutOne native-shaped temporary and its original sharded destination.
_QuantizedProjectionSource [experts, output, input] tensors and a matching floating destination view.

Functions

NameDescription
_expert_amax_storageAllocate [experts] outer-scale metadata, replicated across inner-axis shards.
_localReturn the tensor’s local storage without copying.
_needs_layout_bufferIdentify source-layout operations that intersect a native row shard.
_quantized_storageAllocate source-layout storage for […, input/divisor], preserving shard placements.

Data

_EXPERT_KEY

_PEFT_KEY

API

class nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter(
    config: nemo_automodel.components.models.inkling2.config.Inkling2TextConfig | nemo_automodel.components.models.inkling2.config.Inkling2Config
)

Bases: StateDictAdapter

Map the delivered text weights; preserve floating checkpoint layouts on export.

Multimodal configs additionally map visual parameters without changing their layout. Load destinations normally alias model storage. Quantized experts, row-sharded fused projections and singleton residual queries instead get staging tensors from to_hf(for_checkpoint_load=True); from_hf decodes or copies them into the original storage and reports those keys through view_loaded_native_keys.

_pending
list[tuple[frozenset[str], _QuantizedProjection | _FloatingLayout]] = []
_view_loaded_native_keys
set[str] = set()
multimodal
= isinstance(config, Inkling2Config)
nvfp4_layer_count
source_quantized
= config.quantization_config is not None
view_loaded_native_keys
set[str]

Native keys completed in model storage by from_hf after the most recent DCP load.

The shared loader counts them as loaded although from_hf omits them from its result.

nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter._layout_load_targets(
    name: str,
    tensor: torch.distributed.tensor.DTensor
) -> dict[str, torch.Tensor]

Allocate one replicated native-shaped temporary for a row-sharded layout conversion.

Parameters:

name
str

Native fused projection or singleton residual query name.

tensor
DTensor

Final destination with Shard(0) on one or more mesh axes.

Returns: dict[str, torch.Tensor]

Source-named views of the temporary for DCP to fill.

nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter._quantized_load_targets(
    name: str,
    tensor: torch.Tensor
) -> dict[str, torch.Tensor]

Allocate packed source-layout staging for one grouped expert projection.

Parameters:

name
str

Native grouped expert key, layers.N.mlp.experts.{gate_and_up_projs,down_projs}.

tensor
torch.Tensor

Final floating [experts, input, output] storage; DTensor shards are preserved and Partial placements are rejected.

Returns: dict[str, torch.Tensor]

Uninitialized packed weight, block scales and, for NVFP4 layers, per-expert amax and the

nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter._split_sizes(
    layer: int
) -> tuple[int, ...]
nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    for_checkpoint_load: bool = False,
    kwargs: typing.Any = {}
) -> list[tuple[str, torch.Tensor]]

Return source-layout tensors, gathering row shards only when the layout requires it.

Parameters:

fqn
str

Native text parameter or persistent-buffer name.

tensor
torch.Tensor

Native parameter layout; grouped experts are [experts, input, output].

for_checkpoint_load
boolDefaults to False

Zero omitted fused-projection padding in destination storage before loading. Floating export otherwise leaves the input unchanged.

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

Source names and tensors retaining dtype/device. Fused row splits and

nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter.forced_hf_dtype_mapping(
    state_dict: dict[str, torch.Tensor]
) -> dict[str, str]

Export actual floating dtypes instead of reusing packed source-checkpoint dtype metadata.

Parameters:

state_dict
dict[str, torch.Tensor]

Source-named floating export views, with arbitrary parameter shapes.

Returns: dict[str, str]

SafeTensors dtype names for every exported tensor, including intrinsic FP32 groups.

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

Convert an already-read text checkpoint, preserving interleaving and supplied floating dtypes.

Parameters:

hf_state_dict
dict[str, torch.Tensor]

Delivered source names and floating or packed tensors. After a DCP load into to_hf(for_checkpoint_load=True) destinations, the staged quantized experts and replicated layout temporaries are decoded or copied into model storage here and omitted from the result. model.visual tensors are omitted intentionally by the text-only entry point. Quantized experts read into a full CPU state dict use [E,output,packed_input] storage.

device_mesh
DeviceMesh | NoneDefaults to None

Shared-loader context; DTensor values keep their own placements.

Returns: dict[str, torch.Tensor]

Native tensor mapping. Direct weights are views. Keys completed in model storage are reported

nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter.get_hf_state_dict_keys(
    state_dict: dict[str, torch.Tensor]
) -> list[str]

List released base keys without treating newly injected LoRA rows as base weights.

Parameters:

state_dict
dict[str, torch.Tensor]

Native model tensors, possibly with unprefixed LoRA A/B matrices [rank, input]/[output, rank] or grouped [experts, …, rank].

Returns: list[str]

Released base tensor names only. Adapter saves use the shared PEFT

nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter.to_hf(
    state_dict: dict[str, torch.Tensor],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    for_checkpoint_load: bool = False,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    v4_compatible: bool = False,
    legacy_paramwrapper_layout: bool = False,
    kwargs: typing.Any = {}
) -> dict[str, torch.Tensor]

Export floating source-layout views, or allocate DCP load destinations.

Parameters:

state_dict
dict[str, torch.Tensor]

Native tensors in their original floating layouts and placements.

exclude_key_regex
str | NoneDefaults to None

Optional shared-loader exclusion pattern.

quantization
boolDefaults to False

Allocate packed expert staging matching the released NVFP4/MXFP8 layout. Requires for_checkpoint_load; trained weights are exported without quantization.

for_checkpoint_load
boolDefaults to False

Whether shared checkpointing requested DCP destinations. Row-sharded fused projections and singleton residual queries then get one replicated temporary each, which from_hf redistributes into the original shards.

device_mesh
DeviceMesh | NoneDefaults to None

Shared-loader context; placements are retained on each tensor view.

v4_compatible
boolDefaults to False

Legacy config export is not qualified for this text entry point.

legacy_paramwrapper_layout
boolDefaults to False

Legacy expert LoRA export is unsupported.

Returns: dict[str, torch.Tensor]

Source-named tensors. Export returns floating views and leaves inputs unchanged. Load

class nemo_automodel.components.models.inkling2.state_dict_adapter._FloatingLayout(
    destination: torch.distributed.tensor.DTensor,
    buffer: torch.distributed.tensor.DTensor
)
Dataclass

One native-shaped temporary and its original sharded destination.

buffer
DTensor
destination
DTensor
nemo_automodel.components.models.inkling2.state_dict_adapter._FloatingLayout.finish() -> None

Copy [rows, hidden] or [queries, hidden] into the original DP/EP shard storage.

class nemo_automodel.components.models.inkling2.state_dict_adapter._QuantizedProjection(
    destination: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    amax: torch.Tensor | None = None,
    denominator: torch.Tensor | None = None
)
Dataclass

Source [experts, output, input] tensors and a matching floating destination view.

The destination aliases model storage, potentially a noncontiguous transpose. Packed FP4 has input/2 columns, scales have input/16 or input/32 columns. Only one expert and at most 128 output rows are expanded at a time.

amax
Tensor | None = None
denominator
Tensor | None = None
destination
Tensor
scale
Tensor
weight
Tensor
nemo_automodel.components.models.inkling2.state_dict_adapter._QuantizedProjection.finish() -> None

Decode source tiles in FP32, cast once and copy into the existing destination.

nemo_automodel.components.models.inkling2.state_dict_adapter._expert_amax_storage(
    template: torch.Tensor
) -> torch.Tensor

Allocate [experts] outer-scale metadata, replicated across inner-axis shards.

Parameters:

template
torch.Tensor

Expert tensor [experts, …]; a DTensor supplies the mesh and placements, of which only a leading-axis shard is kept.

Returns: torch.Tensor

Newly allocated, uninitialized FP32 [experts] storage on the template’s device. A

nemo_automodel.components.models.inkling2.state_dict_adapter._local(
    tensor: torch.Tensor
) -> torch.Tensor

Return the tensor’s local storage without copying.

Parameters:

tensor
torch.Tensor

Tensor of shape […], with arbitrary rank and axis order. A DTensor describes the global shape and its mesh placements.

Returns: torch.Tensor

Tensor with the same axis order, dtype and device. For a DTensor, each dimension

nemo_automodel.components.models.inkling2.state_dict_adapter._needs_layout_buffer(
    name: str,
    tensor: torch.Tensor
) -> bool

Identify source-layout operations that intersect a native row shard.

Parameters:

name
str

Native parameter name.

tensor
torch.Tensor

Projection [rows, hidden] or residual query [queries, hidden], optionally DTensor with Shard(0) on one or more mesh axes.

Returns: bool

True for a sharded fused row split or removal of a sharded singleton axis.

nemo_automodel.components.models.inkling2.state_dict_adapter._quantized_storage(
    template: torch.Tensor,
    divisor: int,
    dtype: torch.dtype
) -> torch.Tensor

Allocate source-layout storage for […, input/divisor], preserving shard placements.

Parameters:

template
torch.Tensor

Floating destination […, input]; a DTensor supplies the global shape and placements, and both its global and local trailing extents must be multiples of divisor.

divisor
int

Source elements per storage element along the trailing axis: 2 for packed FP4 pairs, the block size for block scales.

dtype
torch.dtype

Storage dtype of the source tensor.

Returns: torch.Tensor

Newly allocated, uninitialized […, input/divisor] storage on the template’s device.

nemo_automodel.components.models.inkling2.state_dict_adapter._EXPERT_KEY = re.compile('layers\\.(\\d+)\\.mlp\\.experts\\.(?:gate_and_up_projs|down_projs)')
 
nemo_automodel.components.models.inkling2.state_dict_adapter._PEFT_KEY = re.compile('base_model\\.model\\.layers\\.(\\d+)\\.(?:attn\\.(?:qkvr|wo_ud)\\.lo...