ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsDeepseek V41

nemo_automodel.components.models.deepseek_v41.state_dict_adapter

View as Markdown

State dict adapter for DeepSeek V4.1.

The released deepseek-ai/DeepSeek-V4.1-Flash safetensors follow the reference inference module tree. On-disk layout (from the shard headers):

  • FP8 E4M3 projections with float8_e8m0fnu scales over 32x32 blocks (attn.{wq_a,wq_b,wkv,wo_a,wo_b}, attn.indexer.wq_b, ffn.shared_experts.w{1,2,3}, engram.wkv). Initialization uses this fixed 32-column layout rather than DeepSeek V4’s 128x128 blocks.
  • FP4 E2M1 routed experts packed two per int8 with per-row / 32-column e8m0 scales.
  • Engram tables: FP8 E4M3 [rows, 256] with per-row / 32-column e8m0 scales [rows, 8].
  • BF16 / FP32 for everything else (norms, gate, hyper-connection mixers, compressor, indexer keys, attention sink, embeddings, head).

Key mapping (HF -> internal): embed.weight -> model.embed_tokens.weight norm.weight -> model.norm.weight head.weight -> lm_head.weight layers.{i}.attn_norm.weight -> model.layers.{i}.attn_norm.weight layers.{i}.ffn_norm.weight -> model.layers.{i}.ffn_norm.weight layers.{i}.attn.attn_sink -> model.layers.{i}.attn.sinks_param.weight layers.{i}.attn.* -> model.layers.{i}.attn.* (compressor., indexer. keep their names) layers.{i}.ffn.gate.bias -> model.layers.{i}.ffn.gate.e_score_correction_bias layers.{i}.ffn.gate.weight -> model.layers.{i}.ffn.gate.weight layers.{i}.ffn.shared_experts.w1/w3/w2 -> model.layers.{i}.ffn.shared_experts.gate_proj/up_proj/down_proj layers.{i}.ffn.experts.{j}.w1/w3/w2 -> stacked into model.layers.{i}.ffn.experts.gate_and_up_projs / down_projs layers.{i}.hc_attn_{fn,base,scale} -> model.layers.{i}.attn_hc.{fn,base,scale} layers.{i}.hc_ffn_{fn,base,scale} -> model.layers.{i}.ffn_hc.{fn,base,scale} layers.{i}.engram.* -> model.layers.{i}.engram.* layers.{i}.ffn.gate.bias_vl -> model.layers.{i}.ffn.gate.bias_vl vision.* / aligner.* -> model.vision.* / model.aligner.* image_{start,end,newline} -> model.image_{start,end,newline}

The mtp.* DSpark draft is excluded. Eager loading also excludes layers beyond the configured backbone and experts assigned to other ranks before dequantization. Other checkpoint keys are not filtered by optional-tower configuration.

Module Contents

Classes

NameDescription
DeepseekV41StateDictAdapterConvert released V4.1 checkpoint layouts for DCP loading and export.

Functions

NameDescription
_local_offsetsLocate a contiguous DTensor shard without gathering its values.
_native_key-
_released_key-
dequantize_checkpoint_weightDecode released FP8 or packed FP4 weights with bounded FP32 temporaries.

Data

_ENGRAM_EMBED_PATTERN

API

class nemo_automodel.components.models.deepseek_v41.state_dict_adapter.DeepseekV41StateDictAdapter(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert released V4.1 checkpoint layouts for DCP loading and export.

Floating DCP initialization uses shared MoE views and skips rebuilding experts already written into model storage. Quantized load targets use the released V4.1 layout; floating export uses the shared expert splitter.

_engram_rows
_expert_path_segment
str
nemo_automodel.components.models.deepseek_v41.state_dict_adapter.DeepseekV41StateDictAdapter._dequantize(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]

Dequantize paired weights and require scales for retained packed tensors.

Parameters:

state_dict
dict[str, Any]

Mutated released-name mapping. Dense FP8 matrices have shape [rows, columns] with scales [ceil(rows / 32), ceil(columns / 32)]. FP4 experts have shape [rows, columns / 2] with scales [rows, columns / 32]. FP8 Engram tables have shape [rows, channels] with scales [rows, channels / 32]. DTensors retain their global shape and mesh placements, including uneven row owners and inner-axis expert shards. Other tensors retain their registered shapes.

Returns: dict[str, Any]

The same mapping with consumed scale entries and dequantized weights

Raises:

  • ValueError: A packed INT8 or FP8 E4M3 weight has no companion scale.
nemo_automodel.components.models.deepseek_v41.state_dict_adapter.DeepseekV41StateDictAdapter._engram_checkpoint_tensor(
    tensor: torch.Tensor,
    layer_id: int
) -> torch.Tensor

Expose logical checkpoint rows without gathering owner storage.

Parameters:

tensor
torch.Tensor

Table of global shape [padded_rows, channels], either a complete local tensor or a DTensor with placement Shard(0) on a one-dimensional owner mesh. Each owner stores equal rows.

layer_id
int

Decoder layer identifying the logical checkpoint row count.

Returns: torch.Tensor

Aliasing view of shape [rows, channels]. A DTensor preserves its

nemo_automodel.components.models.deepseek_v41.state_dict_adapter.DeepseekV41StateDictAdapter._quantized_load_targets(
    key: str,
    value: torch.Tensor
) -> list[tuple[str, torch.Tensor]]
staticmethod

Allocate rank-local destinations matching the original quantized dump.

Parameters:

key
str

Released checkpoint matrix name.

value
torch.Tensor

Dequantized matrix [rows, columns], possibly a DTensor. Row scales retain row sharding; FP4 column shards must start and end on 32-column block boundaries.

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

Packed INT8 [rows, columns / 2] or FP8 [rows, columns] weight and

nemo_automodel.components.models.deepseek_v41.state_dict_adapter.DeepseekV41StateDictAdapter._restore_engram_padding(
    tensor: torch.Tensor,
    layer_id: int
) -> torch.Tensor

Restore equal owner storage after reading logical checkpoint rows.

Parameters:

tensor
torch.Tensor

Table of global shape [rows, channels], optionally a DTensor with placement Shard(0) on a one-dimensional owner mesh and uneven local shape [local_rows, channels].

layer_id
int

Decoder layer identifying the logical checkpoint row count.

Returns: torch.Tensor

Tensor unchanged for a local table. Distributed tables have global

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

Convert one internal tensor to HF keys, optionally emitting on-disk quantized placeholders.

With quantization=True the placeholders mirror the released layout so DCP can validate shapes / dtypes before the adapter dequantizes on load. These uninitialized targets require for_checkpoint_load=True; trained weights must be exported without quantization.

Parameters:

fqn
str

Internal parameter name.

tensor
Any

Parameter in its model layout. Grouped experts have shape [experts, hidden, 2 * intermediate] or [experts, intermediate, hidden], with optional DTensor sharding of the expert and inner matrix axes. Engram tables have global shape [padded_rows, channels], optionally placement Shard(0) on a one-dimensional owner mesh with equal local row counts. Other tensors retain their arbitrary registered shapes.

**kwargs
AnyDefaults to {}

Checkpoint protocol options, including quantization, exclude_key_regex and for_checkpoint_load.

Returns: list[tuple[str, Any]]

Released-name tensor pairs, with split experts in [output, input]

nemo_automodel.components.models.deepseek_v41.state_dict_adapter.DeepseekV41StateDictAdapter.forced_hf_dtype_mapping(
    state_dict: dict[str, typing.Any]
) -> dict[str, str]

Preserve full-precision parameters when checkpoint export casts weights.

Parameters:

state_dict
dict[str, Any]

Native parameter/buffer tensors with arbitrary registered shapes and layouts. Values are inspected only for their dtype.

Returns: dict[str, str]

Released checkpoint keys that must remain float32, including mHC,

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

Convert the released HF checkpoint to the internal format.

Steps: discard DSpark draft tensors, unconstructed layers and non-local experts before dequantization, restore Engram owner padding, rename, and merge experts not already loaded through views into model storage.

Parameters:

hf_state_dict
dict[str, Any]

Consumed mapping of released-name tensors. Per-expert projections have shape [output, input], with FP4 input columns packed two per byte before dequantization. Engram tables have logical shape [rows, channels] and optionally placement Shard(0) on a one-dimensional owner mesh, with uneven local shape [local_rows, channels]. Other tensors retain the layouts documented in this module.

device_mesh
DeviceMesh | NoneDefaults to None

Optional expert mesh selecting local expert IDs and retaining any inner-axis expert FSDP sharding.

**kwargs
Defaults to {}

Additional checkpoint protocol arguments.

Returns: dict[str, Any]

Internal-name tensors. Engram DTensors have global shape

Raises:

  • ValueError: A retained quantized weight has no scale, a scale has no weight, or multiple released keys map to one native key.
  • RuntimeError: A retained expert layer lacks a required local projection.
nemo_automodel.components.models.deepseek_v41.state_dict_adapter.DeepseekV41StateDictAdapter.get_hf_state_dict_keys(
    state_dict: dict[str, typing.Any]
) -> list[str]

Return global checkpoint names without inspecting owner-local values.

Parameters:

state_dict
dict[str, Any]

Native model mapping, including pre-distribution local Engram parameters and meta tensors of arbitrary shapes.

Returns: list[str]

Rank-independent released names. Grouped expert keys expand over

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

Export native tensors under the released checkpoint’s key names.

Parameters:

state_dict
dict[str, Any]

Native tensor mapping, including grouped expert tensors [experts, hidden, 2 * intermediate] and [experts, intermediate, hidden]. Other values retain their registered shapes. DTensors may shard the expert axis and an inner matrix axis; Engram tables have global shape [padded_rows, channels] with Shard(0) on a one-dimensional owner mesh.

exclude_key_regex
str | NoneDefaults to None

Optional regular expression for excluded HF keys.

quantization
boolDefaults to False

Whether initialization needs packed checkpoint targets.

**kwargs
AnyDefaults to {}

Compatibility options from checkpointing, including for_checkpoint_load for destinations overwritten by DCP.

Returns: dict[str, Any]

Released mapping with split expert matrices [output, input]. DTensor

nemo_automodel.components.models.deepseek_v41.state_dict_adapter._local_offsets(
    tensor: torch.distributed.tensor.DTensor
) -> tuple[int, ...]

Locate a contiguous DTensor shard without gathering its values.

Parameters:

tensor
DTensor

DTensor of arbitrary global shape, with Shard or Replicate placements. Repeated sharding on the same axis is supported.

Returns: tuple[int, ...]

Global offsets of this rank’s local shard along each tensor dimension.

nemo_automodel.components.models.deepseek_v41.state_dict_adapter._native_key(
    key: str
) -> str
 
nemo_automodel.components.models.deepseek_v41.state_dict_adapter._released_key(
    key: str
) -> str
 
nemo_automodel.components.models.deepseek_v41.state_dict_adapter.dequantize_checkpoint_weight(
    weight: torch.Tensor,
    scale: torch.Tensor,
    dtype: torch.dtype = torch.bfloat16,
    rowwise: bool = False
) -> torch.Tensor

Decode released FP8 or packed FP4 weights with bounded FP32 temporaries.

Parameters:

weight
torch.Tensor

FP8 tensor of shape [rows, columns], or packed INT8 tensor of shape [rows, columns / 2]. Packed E2M1 stores the even column in the low nibble and the odd column in the high nibble. A DTensor preserves its global shape and Shard/Replicate placements.

scale
torch.Tensor

Tensor of shape [ceil(rows / 32), ceil(columns / 32)] for dense FP8, or [rows, ceil(columns / 32)] for FP4 and Engram FP8. A plain scale may cover the global matrix or exactly this rank’s blocks; a DTensor scale must cover the weight shard at matching offsets.

dtype
torch.dtypeDefaults to torch.bfloat16

Dequantized floating-point storage dtype.

rowwise
boolDefaults to False

Use per-row scales for FP8 Engram tables. FP4 always uses them.

Returns: torch.Tensor

Independent tensor of shape [rows, columns] in dtype; DTensor

nemo_automodel.components.models.deepseek_v41.state_dict_adapter._ENGRAM_EMBED_PATTERN = re.compile('^layers\\.(\\d+)\\.engram\\.embed\\.weight$')