nemo_automodel.components.models.mistral_large4.model

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

Module Contents

Classes

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

Functions

Name	Description
`_mistral_swiglu`	Preserve Mistral’s BF16 SiLU and product rounding before the down GEMM.

API

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

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:

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.

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, HFCheckpointingMixin, PreTrainedModel, 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

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 | NoneDefaults to 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.TensorDefaults to ()

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 | NoneDefaults to None

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

position_ids

torch.Tensor | NoneDefaults to None

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

attention_mask

torch.Tensor | NoneDefaults to 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 | NoneDefaults to None

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

labels

torch.Tensor | NoneDefaults to 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.TensorDefaults to 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

boolDefaults to False

Return final [batch, sequence, hidden] activations.

cu_seqlens

torch.Tensor | NoneDefaults to 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 | NoneDefaults to None

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

qkv_format

strDefaults to 'bshd'

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

use_cache

boolDefaults to False

Must be false; incremental generation caches are not supported.

_mistral_packed_document_ids

torch.Tensor | NoneDefaults to 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],

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.

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

Return the untied input embedding.

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

Return the untied language head.

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 | NoneDefaults to None

Device on which the caller materialized parameters.

dtype

torch.dtype | NoneDefaults to None

Optional final body storage dtype; sensitive modules remain FP32.

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

intDefaults to 1

Number of PP microbatches along the batch dimension.

Returns: dict[str, ContextParallelSharder]

Sharder mapping using prepare_cp_inputs’ tensor contract. PP keeps

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].

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)

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 | NoneDefaults to 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,

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

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

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

Identify trimmed stages from their actual module ownership.

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;

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

Return original, contiguous decoder indices after the shared splitter.

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.

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,

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 | NoneDefaults to None

Resolved backbone supplied by the generic splitter.

Returns: list[list[str]]

Independent FQN lists with LeRoPE exclusively on the first stage.

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.

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.