> 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.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.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.gpt_oss.state_dict_adapter

## Module Contents

### Classes

| Name                                                                                                            | Description |
| --------------------------------------------------------------------------------------------------------------- | ----------- |
| [`GPTOSSStateDictAdapter`](#nemo_automodel-components-models-gpt_oss-state_dict_adapter-GPTOSSStateDictAdapter) | -           |

### Functions

| Name                                                                                                      | Description                                                                 |
| --------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`_finish_mxfp4_loads`](#nemo_automodel-components-models-gpt_oss-state_dict_adapter-_finish_mxfp4_loads) | Install one decoder layer's packed MXFP4 expert tensors into model storage. |

### Data

[`FP4_VALUES`](#nemo_automodel-components-models-gpt_oss-state_dict_adapter-FP4_VALUES)

[`_DECODER_LAYER_KEY`](#nemo_automodel-components-models-gpt_oss-state_dict_adapter-_DECODER_LAYER_KEY)

[`_MXFP4_EXPERT_SUFFIXES`](#nemo_automodel-components-models-gpt_oss-state_dict_adapter-_MXFP4_EXPERT_SUFFIXES)

### API

```python
class nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter(
    config: transformers.GptOssConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)
```

**Bases:** [StateDictAdapter](/nemo-automodel/nemo_automodel/components/checkpoint/state_dict_adapter#nemo_automodel-components-checkpoint-state_dict_adapter-StateDictAdapter)

**`hf_to_internal_map`**

---

**`internal_to_hf_map`**

---

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter._apply_key_mapping(
    state_dict: dict[str, typing.Any],
    mapping: dict[str, str]
) -> dict[str, typing.Any]
```

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter._convert_moe_packed_tensors(
    blocks,
    scales,
    dtype: torch.dtype = torch.bfloat16,
    rows_per_chunk: int = 32768 * 1024
) -> torch.Tensor
```

Decode packed MXFP4 expert weights into independent floating-point storage.

Source: [https://github.com/huggingface/transformers/blob/869735d37d0f929311ac6611728c482a4414ba8c/src/transformers/integrations/mxfp4.py#L77](https://github.com/huggingface/transformers/blob/869735d37d0f929311ac6611728c482a4414ba8c/src/transformers/integrations/mxfp4.py#L77)

**Parameters:**

**`blocks`**

Uint8 tensor of shape \[experts, output\_features, groups, packed\_bytes], where each group stores
`packed_bytes` pairs of FP4 values, low nibble first. Checkpoints use 16 bytes per group.
DTensor inputs may shard the expert axis on at most one mesh axis, with the expert count divisible
by that axis's size; all other placements must replicate. `scales` must use the same mesh and
placements. These are global shapes; each rank decodes its local expert rows. CPU inputs may move
to CUDA when distributed execution has more than one rank.

---

**`scales`**

Uint8 tensor of shape \[experts, output\_features, groups], with one biased base-2 exponent per
group. The decoded values are multiplied by `2 ** (scales - 127)`. Must share `blocks`' device.

---

**`dtype`** `torch.dtype` — default: torch.bfloat16

Floating-point output dtype.

---

**`rows_per_chunk`** `int` — default: 32768 \* 1024

Maximum number of flattened rank-local \[packed\_bytes] rows decoded per chunk.

---

**Returns:** `torch.Tensor`

Tensor of shape \[experts, input\_features, output\_features], where

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter._dequantize_block_scale_tensors(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]
```

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter._iter_checkpoint_load_parts(
    model_state_dict: dict[str, torch.Tensor],
    expert_model_keys_by_layer: dict[int, list[str]]
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart]
```

Build lazy direct and per-layer MXFP4 load parts.

**Parameters:**

**`model_state_dict`** `dict[str, torch.Tensor]`

Native names mapped to final single-device tensors. Expert tensors have layout
`[experts, input_features, output_features]` and BF16 or FP32 dtype. All direct destinations are
mutated in place by DCP.

---

**`expert_model_keys_by_layer`** `dict[int, list[str]]`

Complete decoder-layer indices mapped to the two native expert projection
names owned by each layer.

---

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter._model_to_hf_key(
    model_key: str
) -> str
```

Map one native GPT-OSS tensor name to its Hugging Face name.

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

Convert a single tensor from native format to HuggingFace format.

**Parameters:**

**`fqn`** `str`

Fully qualified name of the tensor in native format

---

**`tensor`** `Any`

The tensor to convert

---

**`**kwargs`** — default: \{}

Additional arguments for conversion

---

**Returns:** `list[tuple[str, Any]]`

List of (fqn, tensor) tuples in HuggingFace format

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

Convert HF checkpoint to native format in-place.

* Apply key mappings from HF to internal format
* Dequantize block/scale tensors (freeing originals)

Operates in-place on the input dict to avoid allocating a full copy,
reducing peak memory from 2x to \~1x model size.

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter.iter_checkpoint_load_parts(
    model_state_dict: dict[str, torch.Tensor],
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart] | None
```

Load a single-GPU GPT-OSS MXFP4 checkpoint one decoder layer at a time.

Ordinary checkpoint tensors load directly into their final model tensors. Each decoder-layer part allocates
only that layer's packed uint8 expert blocks and scales. After DCP fills them, the part decodes one projection
at a time into the existing BF16 or FP32 model tensor and releases the packed values before advancing. Backend
bookkeeping entries ending in `_extra_state` are not checkpoint tensors and keep their initialized values.

This path intentionally requires a complete, non-distributed decoder. Distributed GPT-OSS loading already
uses rank-local DCP tensors, while some distributed expert backends expose state-dict tensors that do not own
the final parameter storage.

**Parameters:**

**`model_state_dict`** `dict[str, torch.Tensor]`

Native names mapped to final model tensors. Expert projection tensors must have layout
`[experts, input_features, output_features]`, use BF16 or FP32, and own ordinary single-device
storage. Other tensors retain arbitrary model-defined shapes, dtypes, devices, strides, and storage.

---

**`device_mesh`** `Optional[DeviceMesh]` — default: None

Optional distributed mesh. A non-`None` mesh disables this single-device path.

---

**Returns:** `Iterator[CheckpointLoadPart] | None`

One direct-load part for ordinary tensors followed by one bounded temporary-load part per decoder layer,

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

Convert from native model state dict to HuggingFace format.

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter._finish_mxfp4_loads(
    adapter: nemo_automodel.components.models.gpt_oss.state_dict_adapter.GPTOSSStateDictAdapter,
    conversions: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]
) -> None
```

Install one decoder layer's packed MXFP4 expert tensors into model storage.

**Parameters:**

**`adapter`** `GPTOSSStateDictAdapter`

GPT-OSS adapter that decodes the checkpoint's MXFP4 representation.

---

**`conversions`** `tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]`

Tuples of `(target, blocks, scales)`. Each `target` is a final BF16 or FP32 model tensor with
layout `[experts, input_features, output_features]`. `blocks` is a temporary uint8 checkpoint tensor
with layout `[experts, output_features, input_features / 32, 16]`; each byte holds two FP4 values.
`scales` is its temporary uint8 exponent tensor with layout
`[experts, output_features, input_features / 32]`. This function mutates each target in place and does
not retain the temporary tensors.

---

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter.FP4_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, -...
```

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter._DECODER_LAYER_KEY = re.compile('^model\\.layers\\.(\\d+)\\.')
```

```python
nemo_automodel.components.models.gpt_oss.state_dict_adapter._MXFP4_EXPERT_SUFFIXES = ('mlp.experts.gate_and_up_projs', 'mlp.experts.down_projs')
```