> 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.*`),
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.

---

```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
) -> 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, sequence, hash\_heads].

---

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

---

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

---

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

---

**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_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
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
) -> 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.

---

**Returns:** `torch.Tensor`

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

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