> 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.deepseek_v41.model

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.*`) are built
separately by :mod:`nemo_automodel.components.models.deepseek_v41.dspark` so
their objective cannot backpropagate into this backbone. Inference-time KV
caching and SWA bounded replay remain out of scope.

## Module Contents

### Classes

| Name                                                                                                    | Description                                                                  |
| ------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`DeepseekV41Block`](#nemo_automodel-components-models-deepseek_v41-model-DeepseekV41Block)             | CSA2 and MoE sublayers with the single-pass mHC coefficient handoff.         |
| [`DeepseekV41ForCausalLM`](#nemo_automodel-components-models-deepseek_v41-model-DeepseekV41ForCausalLM) | DeepSeek V4.1 causal LM with optional vision and an fp32 `lm_head`.          |
| [`DeepseekV41Model`](#nemo_automodel-components-models-deepseek_v41-model-DeepseekV41Model)             | DeepSeek V4.1 decoder stack: embeddings, hyper-connected blocks, final norm. |

### Data

[`ModelClass`](#nemo_automodel-components-models-deepseek_v41-model-ModelClass)

### API

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

---

**`self_attn`** `DeepseekV41Attention`

Expose the attention module to the shared CP parallelizer.

---

```python
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,
    cp_group: torch.distributed.ProcessGroup | None = None,
    packed_seq_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 | None` — default: None

Optional binary tensor of shape \[batch, sequence].

---

**`image_mask`** `torch.Tensor | None` — default: None

Optional boolean tensor of shape \[batch, sequence].

---

**`engram_hash_ids`** `torch.Tensor | None` — default: None

Optional logical memory rows \[batch, local\_sequence, hash\_heads].

---

**`cp_group`** `dist.ProcessGroup | None` — default: None

CP group; all token axes are local and CSA2 keys are global.

---

**`packed_seq_ids`** `torch.Tensor | None` — default: None

Optional document IDs \[batch, local\_sequence], zero for padding.

---

**Returns:** `torch.Tensor`

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

```python
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](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `PreTrainedModel`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-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`**

---

**`_owns_cp_attention`** `bool = True`

---

**`_packed_alignment`** `= math.lcm(*(r for r in ratios if r))`

---

**`backend`**

---

**`base_model_prefix`** `str = 'model'`

---

**`config_class`** `type[DeepseekV41Config] = DeepseekV41Config`

---

**`cp_mesh`** `DeviceMesh | None = None`

---

**`lm_head`**

---

**`model`**

---

**`state_dict_adapter`**

---

**`tie_word_embeddings_support`** `TieSupport = TieSupport.UNTIED_ONLY`

---

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

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

```python
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,
    cp_group: torch.distributed.ProcessGroup | None = None,
    packed_seq_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    seq_lens_padded: torch.Tensor | None = None,
    qkv_format: str | None = None
) -> 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 | None` — default: None

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

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional integer tensor of shape \[batch, sequence].

---

**`labels`** `torch.Tensor | None` — default: None

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

---

**`pixel_values`** `torch.Tensor | None` — default: None

Optional image patches \[all\_patches, 3, patch\_size, patch\_size].

---

**`image_grid_hws`** `torch.Tensor | None` — default: None

Optional patch-grid sizes \[images, 2].

---

**`vision_token_types`** `torch.Tensor | None` — default: None

Optional image/text markers \[batch, sequence].

---

**`logits_to_keep`** `int | torch.Tensor` — default: 0

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

---

**`return_hidden_states`** `bool` — default: False

Return final hidden states for the recipe's loss.

---

**`output_hidden_states`** `bool` — default: False

Capture residual streams for numerical comparisons.

---

**`seq_lens`** `torch.Tensor | None` — default: None

Optional real document lengths \[batch, documents] for packed text.

---

**`seq_lens_padded`** `torch.Tensor | None` — default: None

Optional physical document spans \[batch, documents].

---

**`qkv_format`** `str | None` — default: None

Optional "thd" marker from packed\_sequence\_thd\_collater.

---

**`packed_seq_ids`** `torch.Tensor | None` — default: None

Prepared document IDs \[batch, local\_sequence], zero for padding.
The model-owned CP sharder supplies these after aligning and sharding a pack.

---

**`cp_group`** `dist.ProcessGroup | None` — default: None

CP group; input/output token axes contain only the local
contiguous shard. The recipe computes loss from globally shifted labels.

---

**Returns:** `CausalLMOutputWithPast`

CausalLMOutputWithPast containing logits \[batch, kept\_sequence, vocab],

```python
nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM.from_config(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41Config,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM
```

classmethod

Construct using the NeMo registry's configuration entry point.

```python
nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM.get_dspark_target_feature_modules(
    layer_ids: list[int]
) -> tuple[torch.nn.Module, ...]
```

Return modules whose inputs are the released DSpark target features.

The reference implementation captures the residual streams immediately
before attention in each selected layer, after that layer's optional
Engram update. The attention hyper-connection is the first module to
consume those streams, so its forward input is the exact capture point.

**Parameters:**

**`layer_ids`** `list[int]`

Strictly increasing decoder-layer indices in
`[0, num_hidden_layers)`.

---

**Returns:** `nn.Module`

Modules ordered like `layer_ids`. Each receives a tensor of shape

**Raises:**

* `ValueError`: If the indices are duplicated, unsorted, or outside the
  active decoder depth.

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

Return the untied token embedding module.

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

Return the independent vocabulary projection.

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

```python
nemo_automodel.components.models.deepseek_v41.model.DeepseekV41ForCausalLM.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]
```

Select contiguous text sharding with complete compression groups.

**Parameters:**

**`batch`** `dict[str, Any]`

Full text tensors of shape \[batch, global\_sequence]. The hook
leaves tensors intact; sharding happens through the returned strategy.

---

**`num_chunks`** `int` — default: 1

Framework chunk count; this model uses one chunk.

---

**Returns:** `dict[str, Any]`

Model-owned CP sharder. Its local token tensors have shape

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

---

```python
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,
    cp_group: torch.distributed.ProcessGroup | None = None,
    packed_seq_ids: torch.Tensor | None = None
) -> 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 | None` — default: None

Optional integer tensor of shape \[batch, sequence].

---

**`attention_mask`** `torch.Tensor | None` — default: None

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

---

**`image_mask`** `torch.Tensor | None` — default: None

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

---

**`inputs_embeds`** `torch.Tensor | None` — default: None

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

---

**`output_hidden_states`** `bool` — default: False

Whether to retain streams before each block.

---

**`packed_seq_ids`** `torch.Tensor | None` — default: None

Optional document IDs \[batch, local\_sequence], zero for padding.

---

**`cp_group`** `dist.ProcessGroup | None` — default: None

CP group; token axes above are local sequence shards.
Engram hashes use global token history, with only local rows looked up.

---

**Returns:** `torch.Tensor`

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

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