ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsDeepseek V41

nemo_automodel.components.models.deepseek_v41.model

View as Markdown

DeepSeek V4.1 text and image backbone for AutoModel training.

Forward contract (reference Transformer.forward of inference/model.py):

  1. Embed tokens and expand the hidden state into hc_mult residual streams.
  2. Run the 40 blocks with single-pass mHC: each block receives the input mix produced by the previous block’s FFN site (a one-hot mix reads stream 0 at the start). Engram modules write into the streams before their layer.
  3. Collapse the streams with the last FFN-site mix, apply the final RMSNorm and the fp32 lm_head.

Cross-layer CSA2 state (shared compressed KV, index keys, Top-K indices and the hierarchical candidate pool) lives in per-layer snapshots of :class:~nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState. Snapshots share tensors and preserve the state needed for activation recomputation.

The optional vision tower inserts projected image patches and learned image delimiters into the text sequence. Text and image batches use full sequences with two-dimensional token layouts. DSpark draft layers (mtp.*), inference-time KV caching, and SWA bounded replay remain out of scope.

Module Contents

Classes

NameDescription
DeepseekV41BlockCSA2 and MoE sublayers with the single-pass mHC coefficient handoff.
DeepseekV41ForCausalLMDeepSeek V4.1 causal LM with optional vision and an fp32 lm_head.
DeepseekV41ModelDeepSeek V4.1 decoder stack: embeddings, hyper-connected blocks, final norm.

Data

ModelClass

API

class nemo_automodel.components.models.deepseek_v41.model.DeepseekV41Block(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    engram_process_group: torch.distributed.ProcessGroup | None = None
)

Bases: Module

CSA2 and MoE sublayers with the single-pass mHC coefficient handoff.

attn
= DeepseekV41Attention(config, layer_idx, backend)
attn_hc
attn_norm
engram
ffn
= MoE(moe_config, backend)
ffn_hc
ffn_norm
mlp
MoE

Expose the shared parallelizer’s MoE interface without duplicate registration.

nemo_automodel.components.models.deepseek_v41.model.DeepseekV41Block.forward(
    hidden_states: torch.Tensor,
    pre_mix: torch.Tensor,
    state: nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState,
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    image_mask: torch.Tensor | None = None,
    engram_hash_ids: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor, nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState]

Execute one block while retaining differentiable shared KV ownership.

Parameters:

hidden_states
torch.Tensor

Tensor of shape [batch, sequence, streams, hidden].

pre_mix
torch.Tensor

FP32 tensor of shape [batch, sequence, streams].

state
DeepseekV41AttentionState

Shared CSA2 tensors with layouts documented by DeepseekV41AttentionState; no tensor is mutated.

position_ids
torch.Tensor

Integer tensor of shape [batch, sequence].

attention_mask
torch.Tensor | NoneDefaults to None

Optional binary tensor of shape [batch, sequence].

image_mask
torch.Tensor | NoneDefaults to None

Optional boolean tensor of shape [batch, sequence].

engram_hash_ids
torch.Tensor | NoneDefaults to None

Optional logical memory rows [batch, sequence, hash_heads].

Returns: torch.Tensor

Updated streams [batch, sequence, streams, hidden], the next pre-mix

class nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    tokenizer: transformers.PreTrainedTokenizerFast | None = None,
    engram_process_group: torch.distributed.ProcessGroup | None = None
)

Bases: HFCheckpointingMixin, PreTrainedModel, MoEFSDPSyncMixin

DeepSeek V4.1 causal LM with optional vision and an fp32 lm_head.

engram_process_group explicitly selects contiguous row owners for the Engram tables. Distributed models default to WORLD, including a one-rank WORLD. Without distributed initialization, tables remain local. FSDP’s shard mesh must match the owner group exactly.

_keep_in_fp32_modules_strict
backend
base_model_prefix
str = 'model'
config_class
type[DeepseekV41Config] = DeepseekV41Config
lm_head
model
state_dict_adapter
tie_word_embeddings_support
TieSupport = TieSupport.UNTIED_ONLY
nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM._image_embeddings(
    input_ids: torch.Tensor,
    pixel_values: torch.Tensor,
    image_grid_hws: torch.Tensor,
    vision_token_types: torch.Tensor
) -> torch.Tensor

Insert image features and learned delimiters into fresh token embeddings.

Parameters:

input_ids
torch.Tensor

Integer tensor of shape [batch, sequence].

pixel_values
torch.Tensor

Tensor of shape [all_patches, 3, patch_size, patch_size].

image_grid_hws
torch.Tensor

Integer tensor of shape [images, 2], containing patch grids.

vision_token_types
torch.Tensor

Integer tensor of shape [batch, sequence], with -1 for text and 0/1/2/3 for image start/content/newline/end.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden], retaining text and image

nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM._nemo_prepare_model_owned_dtensors(
    fsdp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> set[torch.nn.Parameter]

Register owner table DTensors before FSDP records ignored parameters.

Parameters:

fsdp_mesh
DeviceMesh

One-dimensional shard mesh whose ranks and ordering must match the Engram owner group.

Returns: set[nn.Parameter]

Exact registered parameter identities to exclude from FSDP. Each

nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    image_grid_hws: torch.Tensor | None = None,
    vision_token_types: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    return_hidden_states: bool = False,
    output_hidden_states: bool = False
) -> transformers.modeling_outputs.CausalLMOutputWithPast

Compute full-vocabulary logits or hidden states for the training loss.

Parameters:

input_ids
torch.Tensor

Integer tensor of shape [batch, sequence].

attention_mask
torch.Tensor | NoneDefaults to None

Optional binary right-padding tensor [batch, sequence].

position_ids
torch.Tensor | NoneDefaults to None

Optional integer tensor of shape [batch, sequence].

labels
torch.Tensor | NoneDefaults to None

Optional targets of shape [batch, sequence], with -100 ignored.

pixel_values
torch.Tensor | NoneDefaults to None

Optional image patches [all_patches, 3, patch_size, patch_size].

image_grid_hws
torch.Tensor | NoneDefaults to None

Optional patch-grid sizes [images, 2].

vision_token_types
torch.Tensor | NoneDefaults to None

Optional image/text markers [batch, sequence].

logits_to_keep
int | torch.TensorDefaults to 0

Number of final positions, or integer position indices [kept].

return_hidden_states
boolDefaults to False

Return final hidden states for the recipe’s loss.

output_hidden_states
boolDefaults to False

Capture residual streams for numerical comparisons.

Returns: CausalLMOutputWithPast

CausalLMOutputWithPast containing logits [batch, kept_sequence, vocab],

classmethod

Construct using the NeMo registry’s configuration entry point.

nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM.get_input_embeddings() -> torch.nn.Embedding

Return the untied token embedding module.

nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM.get_output_embeddings() -> torch.nn.Module

Return the independent vocabulary projection.

nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

Initialize every trainable backbone weight after meta materialization.

class nemo_automodel.components.models.deepseek_v41.model.DeepseekV41Model(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    tokenizer: transformers.PreTrainedTokenizerFast | None = None,
    engram_process_group: torch.distributed.ProcessGroup | None = None
)

Bases: Module

DeepSeek V4.1 decoder stack: embeddings, hyper-connected blocks, final norm.

embed_tokens
engram_hash
layers
norm
nemo_automodel.components.models.deepseek_v41.model.DeepseekV41Model.forward(
    input_ids: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    image_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool = False
) -> tuple[torch.Tensor, tuple[torch.Tensor, ...] | None]

Compute all positions without inference-only prefill shortcuts.

Parameters:

input_ids
torch.Tensor

Integer tensor of shape [batch, sequence], also used for Engram.

position_ids
torch.Tensor | NoneDefaults to None

Optional integer tensor of shape [batch, sequence].

attention_mask
torch.Tensor | NoneDefaults to None

Optional binary right-padding tensor [batch, sequence].

image_mask
torch.Tensor | NoneDefaults to None

Optional boolean image-span tensor [batch, sequence].

inputs_embeds
torch.Tensor | NoneDefaults to None

Optional projected multimodal embeddings [batch, sequence, hidden].

output_hidden_states
boolDefaults to False

Whether to retain streams before each block.

Returns: torch.Tensor

Final normalized hidden states [batch, sequence, hidden] and optional

nemo_automodel.components.models.deepseek_v41.model.ModelClass = DeepseekV41ForCausalLM