nemo_automodel.components.models.gpt_oss.state_dict_adapter

Module Contents

Classes

Name	Description
`GPTOSSStateDictAdapter`	-

Functions

Name	Description
`_finish_mxfp4_loads`	Install one decoder layer’s packed MXFP4 expert tensors into model storage.

Data

FP4_VALUES

_DECODER_LAYER_KEY

_MXFP4_EXPERT_SUFFIXES

API

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

hf_to_internal_map

internal_to_hf_map

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]

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

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.dtypeDefaults to torch.bfloat16

Floating-point output dtype.

rows_per_chunk

intDefaults to 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

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

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.

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.

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

Defaults to {}

Additional arguments for conversion

Returns: list[tuple[str, Any]]

List of (fqn, tensor) tuples in HuggingFace format

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.

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]Defaults to 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,

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.

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.

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, -...

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

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