> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.

# nemo_automodel.components.models.inkling2.state_dict_adapter

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

## Module Contents

### Classes

| Name                                                                                                                 | Description                                                                        |
| -------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- |
| [`Inkling2StateDictAdapter`](#nemo_automodel-components-models-inkling2-state_dict_adapter-Inkling2StateDictAdapter) | Map the delivered text weights; preserve floating checkpoint layouts on export.    |
| [`_FloatingLayout`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_FloatingLayout)                   | One native-shaped temporary and its original sharded destination.                  |
| [`_QuantizedProjection`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_QuantizedProjection)         | Source \[experts, output, input] tensors and a matching floating destination view. |

### Functions

| Name                                                                                                         | Description                                                                            |
| ------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------- |
| [`_expert_amax_storage`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_expert_amax_storage) | Allocate \[experts] outer-scale metadata, replicated across inner-axis shards.         |
| [`_local`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_local)                             | Return the tensor's local storage without copying.                                     |
| [`_needs_layout_buffer`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_needs_layout_buffer) | Identify source-layout operations that intersect a native row shard.                   |
| [`_quantized_storage`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_quantized_storage)     | Allocate source-layout storage for \[..., input/divisor], preserving shard placements. |

### Data

[`_EXPERT_KEY`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_EXPERT_KEY)

[`_PEFT_KEY`](#nemo_automodel-components-models-inkling2-state_dict_adapter-_PEFT_KEY)

### API

```python
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](/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/checkpoint/state_dict_adapter#nemo_automodel-components-checkpoint-state_dict_adapter-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.

---

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

```python
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

```python
nemo_automodel.components.models.inkling2.state_dict_adapter.Inkling2StateDictAdapter._split_sizes(
    layer: int
) -> tuple[int, ...]
```

```python
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`** `bool` — default: 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

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

```python
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 | None` — default: 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

```python
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

```python
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 | None` — default: None

Optional shared-loader exclusion pattern.

---

**`quantization`** `bool` — default: False

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

---

**`for_checkpoint_load`** `bool` — default: 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 | None` — default: None

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

---

**`v4_compatible`** `bool` — default: False

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

---

**`legacy_paramwrapper_layout`** `bool` — default: False

Legacy expert LoRA export is unsupported.

---

**Returns:** `dict[str, torch.Tensor]`

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

```python
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`

---

```python
nemo_automodel.components.models.inkling2.state_dict_adapter._FloatingLayout.finish() -> None
```

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

```python
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`

---

```python
nemo_automodel.components.models.inkling2.state_dict_adapter._QuantizedProjection.finish() -> None
```

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

```python
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

```python
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

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

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

```python
nemo_automodel.components.models.inkling2.state_dict_adapter._EXPERT_KEY = re.compile('layers\\.(\\d+)\\.mlp\\.experts\\.(?:gate_and_up_projs|down_projs)')
```

```python
nemo_automodel.components.models.inkling2.state_dict_adapter._PEFT_KEY = re.compile('base_model\\.model\\.layers\\.(\\d+)\\.(?:attn\\.(?:qkvr|wo_ud)\\.lo...
```