> This page is for version Nightly (default).
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

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

Mistral Large 4 text model with expert, pipeline and context parallelism.

## Module Contents

### Classes

| Name                                                                                                              | Description                                                                     |
| ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`MistralLarge4Block`](#nemo_automodel-components-models-mistral_large4-model-MistralLarge4Block)                 | Pre-normalized MLA followed by dense or routed/shared SwiGLU.                   |
| [`MistralLarge4ForCausalLM`](#nemo_automodel-components-models-mistral_large4-model-MistralLarge4ForCausalLM)     | Text model; source FP8 weights are reconstructed on load.                       |
| [`MistralLarge4LayerState`](#nemo_automodel-components-models-mistral_large4-model-MistralLarge4LayerState)       | Per-forward state; never retained on a module or detached.                      |
| [`MistralLarge4Model`](#nemo_automodel-components-models-mistral_large4-model-MistralLarge4Model)                 | Text backbone; each forward owns the lifetime of its shared source activations. |
| [`MistralLarge4PipelineMixin`](#nemo_automodel-components-models-mistral_large4-model-MistralLarge4PipelineMixin) | Use the existing PP hooks and tuple transport without replicated parameters.    |

### Functions

| Name                                                                                        | Description                                                             |
| ------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_mistral_swiglu`](#nemo_automodel-components-models-mistral_large4-model-_mistral_swiglu) | Preserve Mistral's BF16 SiLU and product rounding before the down GEMM. |

### API

```python
class nemo_automodel.components.models.mistral_large4.model.MistralLarge4Block(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    layer_idx: int,
    dtype: torch.dtype
)
```

**Bases:** `Module`

Pre-normalized MLA followed by dense or routed/shared SwiGLU.

**`input_layernorm`**

---

**`mlp`**

---

**`moe_config`**

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4Block.forward(
    x: torch.Tensor,
    tables: tuple[torch.Tensor, torch.Tensor],
    layout: nemo_automodel.components.models.mistral_large4.attention.SequenceLayout,
    shared_kv: tuple[torch.Tensor, torch.Tensor] | None,
    indices: torch.Tensor | None
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor], torch.Tensor]
```

Run one block with immutable sharing inputs, safe for recomputation.

**Parameters:**

**`x`** `torch.Tensor`

Unsharded \[tokens, hidden] activations.

---

**`tables`** `tuple[torch.Tensor, torch.Tensor]`

FP32 cosine/sine pair \[tokens, rotary_dim / 2].

---

**`layout`** `SequenceLayout`

Positions, document IDs and validity \[tokens].

---

**`shared_kv`** `tuple[torch.Tensor, torch.Tensor] | None`

Optional latent \[tokens, latent_dim], rotated keys
\[tokens, rotary_dim] (zero final dimension for NoPE).

---

**`indices`** `torch.Tensor | None`

Optional preceding int64 \[tokens, min(topk, tokens)] selection.

---

**Returns:** `torch.Tensor`

Hidden \[tokens, hidden], KV pair and selected indices in input layouts.

```python
class nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)
```

**Bases:** [MistralLarge4PipelineMixin](#nemo_automodel-components-models-mistral_large4-model-MistralLarge4PipelineMixin), [HFCheckpointingMixin](/nemo/automodel/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/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

Text model; source FP8 weights are reconstructed on load.

Training runs in floating point: serving FP8 indexer/KV-cache execution and
DSpark drafting are not implemented. DSA parameters remain frozen pending a
training specification.

**`_keep_in_fp32_modules_strict`**

---

**`_no_split_modules`** `= ['MistralLarge4Block']`

---

**`_packed_cp_attn_backends`** `= ('tilelang',)`

---

**`backend`**

---

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

---

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

---

**`lm_head`**

---

**`model`**

---

**`moe_config`**

---

**`packed_mask_type`** `= 'document_ids'`

---

**`parallel_layer_groups`** `= {'language': ('model.layers',)}`

---

**`parallelizer`** `= MistralLarge4Parallelizer()`

---

**`state_dict_adapter`**

---

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

---

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    pipeline_tensors: torch.Tensor = (),
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool = False,
    cu_seqlens: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    qkv_format: str = 'bshd',
    use_cache: bool = False,
    _mistral_packed_document_ids: torch.Tensor | None = None
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor | tuple[torch.Tensor, ...]
```

Run causal or document-packed text with capped logits and optional shifted CE.

**Parameters:**

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

Integer \[batch, sequence], or packed \[tokens] with thd
format. Under CP, text-only PP1 takes pre-sharded local IDs;
stage zero of PP takes global IDs and shards after embedding.

---

**`pipeline_tensors`** `torch.Tensor` — default: ()

Non-first PP boundary tensors in the layouts
documented by `_pipeline_boundary_metas`. input_ids then holds
compute-dtype residual \[batch, sequence, hidden].

---

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

Floating \[batch, sequence, hidden] or packed
\[tokens, hidden], replacing IDs with the same CP ownership.

---

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

Nonnegative integer positions \[batch, local_sequence]
under CP, otherwise matching the input token axes.

---

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

Binary keep mask or nonnegative packed document IDs,
matching input token axes. CP stage zero receives a global mask
and shards it alongside embeddings; later stages use transported layout.

---

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

Optional bool padding mask; true excludes a token.
CP sharder supplies \[batch, local_sequence], including CP padding.

---

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

Optional integer targets, same token shape; -100 ignored. This
HF-style path shifts labels once, excluding padding/document edges.

---

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

Zero for all, positive tail length, or int64 \[kept_positions].
Labels require all positions. Index selection applies to the sequence axis.

---

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

Return final \[batch, sequence, hidden] activations.

---

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

SDPA THD requires integer offsets \[documents + 1]
strictly partitioning flattened storage. Prepared THD receives padded
global offsets \[1, documents + 1] from its sharder; its local
document IDs and padding mask preserve inter-document padding.

---

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

Optional real document lengths \[batch, documents], with
-1000 padding, retained by the packed sharder for document-count validation.

---

**`qkv_format`** `str` — default: 'bshd'

bshd or thd (packed output has a leading batch dimension of one).

---

**`use_cache`** `bool` — default: False

Must be false; incremental generation caches are not supported.

---

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

Model-prepared integer \[batch, local_sequence]
document identity for prepared THD. Each microbatch contains one packed row.

---

**Returns:** `CausalLMOutputWithPast | torch.Tensor | tuple[torch.Tensor, ...]`

CausalLMOutputWithPast with capped logits \[batch, kept_sequence, vocab],

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM.from_config(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
) -> nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM
```

classmethod

Construct the registered model through AutoModel's existing interface.

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM.get_input_embeddings() -> torch.nn.Embedding
```

Return the untied input embedding.

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM.get_output_embeddings() -> torch.nn.Linear
```

Return the untied language head.

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype | None = None
) -> None
```

Initialize materialized parameters in their resident storage dtype.

Initialize in place without temporarily promoting the entire pretrained
base to FP32. Protected modules own separate FP32
FSDP units and must never pass through the broad body-dtype cast.

**Parameters:**

**`buffer_device`** `torch.device | None` — default: None

Device on which the caller materialized parameters.

---

**`dtype`** `torch.dtype | None` — default: None

Optional final body storage dtype; sensitive modules remain FP32.

---

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4ForCausalLM.prepare_model_inputs_for_cp(
    batch: nemo_automodel.components.models.mistral_large4.context_parallel.CPBatch,
    num_chunks: int = 1
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]
```

Prepare CP-local targets and positions with the model-owned sharder.

**Parameters:**

**`batch`** `CPBatch`

Global IDs/masks \[batch, sequence] or embeddings
\[batch, sequence, hidden], with optional complete image payloads.
THD requires IDs \[microbatches, packed_sequence] plus
seq_lens/seq_lens_padded \[microbatches, documents].

---

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

Number of PP microbatches along the batch dimension.

---

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

Sharder mapping using prepare_cp_inputs' tensor contract. PP keeps

```python
class nemo_automodel.components.models.mistral_large4.model.MistralLarge4LayerState(
    tables: tuple[torch.Tensor, torch.Tensor],
    shared_kv: dict[int, tuple[torch.Tensor, torch.Tensor]],
    indices: torch.Tensor | None = None
)
```

Dataclass

Per-forward state; never retained on a module or detached.

**`indices`** `Tensor | None = None`

Optional int64 \[tokens, min(index_topk, tokens)] selected keys.

---

**`shared_kv`** `dict[int, tuple[Tensor, Tensor]]`

Source-layer IDs mapped to latent \[tokens, latent_dim] and
rotated key \[tokens, rotary_dim] tensors; NoPE keys have width zero.

---

**`tables`** `tuple[Tensor, Tensor]`

FP32 cosine/sine tensors, each \[tokens, rotary_dim / 2].

---

```python
class nemo_automodel.components.models.mistral_large4.model.MistralLarge4Model(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)
```

**Bases:** `Module`

Text backbone; each forward owns the lifetime of its shared source activations.

**`_last_consumers`**

---

**`embed_tokens`**

---

**`layers`**

---

**`norm`**

---

**`rotary_emb`** `= MistralLarge4Rotary(config)`

---

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4Model.forward(
    hidden: torch.Tensor,
    layout: nemo_automodel.components.models.mistral_large4.attention.SequenceLayout,
    state: nemo_automodel.components.models.mistral_large4.model.MistralLarge4LayerState | None = None
) -> tuple[torch.Tensor, nemo_automodel.components.models.mistral_large4.model.MistralLarge4LayerState]
```

Run owned layers with explicit, per-microbatch shared state.

**Parameters:**

**`hidden`** `torch.Tensor`

Unsharded compute-dtype activations \[tokens, hidden].

---

**`layout`** `SequenceLayout`

Positions, document IDs and validity, each \[tokens].

---

**`state`** `MistralLarge4LayerState | None` — default: None

Optional received FP32 tables \[tokens, rotary_dim / 2], KV
pairs \[tokens, latent_dim] / \[tokens, rotary_width], and int64
selections \[tokens, min(index_topk, tokens)]. Never mutated.

---

**Returns:** `torch.Tensor`

Activations \[tokens, hidden], normalized only on the final stage,

```python
class nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin()
```

Use the existing PP hooks and tuple transport without replicated parameters.

**`_pp_keep_self_forward`** `bool = True`

---

**`_pp_return_hidden_states`** `bool = False`

---

**`_pp_return_hidden_states_supported`** `bool = True`

---

**`config`** `MistralLarge4Config`

---

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

---

**`lm_head`** `Linear | None`

---

**`model`** `Module`

---

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin._finish_pipeline_stage(
    hidden: torch.Tensor,
    layout: nemo_automodel.components.models.mistral_large4.attention.SequenceLayout,
    state: nemo_automodel.components.models.mistral_large4.model.MistralLarge4LayerState
) -> torch.Tensor | tuple[torch.Tensor, ...]
```

Emit capped final logits or fresh tensors for the next stage.

**Parameters:**

**`hidden`** `torch.Tensor`

Compute-dtype residual \[batch, sequence, hidden].

---

**`layout`** `SequenceLayout`

Flattened integer positions/documents and bool validity \[tokens].

---

**`state`** `MistralLarge4LayerState`

FP32 cosine/sine \[tokens, rotary_dim / 2], live KV pairs
\[tokens, latent_dim] / \[tokens, rotary_width], and selected
int64 keys \[tokens, min(index_topk, tokens)].

---

**Returns:** `torch.Tensor | tuple[torch.Tensor, ...]`

Final FP32 logits \[batch, sequence, vocab], or compute-dtype hidden

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin._is_pipeline_parallel_stage() -> bool
```

Identify trimmed stages from their actual module ownership.

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin._pipeline_boundary_metas(
    boundary: int,
    batch: int,
    sequence: int,
    dtype: torch.dtype
) -> tuple[torch.Tensor, ...]
```

Describe residual, rotary, layout and live sharing tensors at one cut.

**Returns:** `torch.Tensor`

Meta tensors: residual \[batch, sequence, hidden] in compute dtype;

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin._pipeline_layer_ids() -> list[int]
```

Return original, contiguous decoder indices after the shared splitter.

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin._pipeline_live_sources(
    boundary: int
) -> list[int]
```

Select produced KV sources still needed at or after a decoder boundary.

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin._receive_pipeline_state(
    hidden: torch.Tensor,
    tensors: tuple[torch.Tensor, ...]
) -> tuple[nemo_automodel.components.models.mistral_large4.attention.SequenceLayout, nemo_automodel.components.models.mistral_large4.model.MistralLarge4LayerState]
```

Decode the prior stage's tensors without detaching differentiable state.

**Parameters:**

**`hidden`** `torch.Tensor`

Compute-dtype residual \[batch, sequence, hidden].

---

**`tensors`** `tuple[torch.Tensor, ...]`

Remaining boundary tensors in `_pipeline_boundary_metas`
order; no tensor aliases are mutated.

---

**Returns:** `SequenceLayout`

Layout with flattened \[tokens] position/document/validity fields,

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]
```

Assign the single learnable LeRoPE module to stage zero.

**Parameters:**

**`module_names_per_stage`** `list[list[str]]`

Generated module FQNs in forward stage order.

---

**`layers_prefix`** `str`

Decoder parent prefix, required to be `model.`.

---

**`text_model`** `nn.Module | None` — default: None

Resolved backbone supplied by the generic splitter.

---

**Returns:** `list[list[str]]`

Independent FQN lists with LeRoPE exclusively on the first stage.

```python
nemo_automodel.components.models.mistral_large4.model.MistralLarge4PipelineMixin.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]
```

Provide exact tuple metadata to the existing PyTorch PP schedule.

**Parameters:**

**`is_first`** `bool`

Whether this is the first logical stage.

---

**`microbatch_size`** `int`

Number of documents in one pipeline microbatch.

---

**`seq_len`** `int`

Global nonpacked sequence length before CP padding.

---

**`dtype`** `torch.dtype`

Body compute dtype; rotary transport always remains FP32.

---

**Returns:** `tuple[torch.Tensor, ...]`

Input/output meta tuples using `_pipeline_boundary_metas`' layouts.

```python
nemo_automodel.components.models.mistral_large4.model._mistral_swiglu(
    projections: torch.Tensor,
    weights: torch.Tensor
) -> torch.Tensor
```

Preserve Mistral's BF16 SiLU and product rounding before the down GEMM.

**Parameters:**

**`projections`** `torch.Tensor`

BF16 or FP32 tensor \[routed_tokens, 2\*intermediate], gate
then up. May be noncontiguous. No input is modified.

---

**`weights`** `torch.Tensor`

All-one activation weights \[routed_tokens, 1]; actual FP32
routing weights are applied by the expert backend after down.

---

**Returns:** `torch.Tensor`

Independent tensor \[routed_tokens, intermediate] in projections.dtype.