nemo_automodel.components.models.mistral_large4.rope_utils

Mistral Large 4 differentiable rotary and SSMax floating-point numerics.

Module Contents

Classes

Name	Description
`FP32LayerNorm`	Normalize indexer keys with an FP32 parameter-owning module call.
`MistralLarge4Rotary`	Shared learnable LeRoPE frequencies; build angles in FP32 after dtype casts.
`SSMax`	Per-head s * log(position + 1) + b in the query’s compute dtype.

Functions

Name	Description
`apply_interleaved_rope`	Rotate adjacent channel pairs in FP32, then restore the activation dtype.

API

class nemo_automodel.components.models.mistral_large4.rope_utils.FP32LayerNorm()

Bases: LayerNorm

Normalize indexer keys with an FP32 parameter-owning module call.

nemo_automodel.components.models.mistral_large4.rope_utils.FP32LayerNorm.forward(
    x: torch.Tensor
) -> torch.Tensor

Normalize in FP32 and restore the input activation dtype.

Parameters:

torch.Tensor

Tensor of shape […, index_dim], with arbitrary leading dimensions.

Returns: torch.Tensor

Independent tensor with the input shape/dtype. Parameters remain FP32.

class nemo_automodel.components.models.mistral_large4.rope_utils.MistralLarge4Rotary(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config
)

Bases: Module

Shared learnable LeRoPE frequencies; build angles in FP32 after dtype casts.

alpha

nemo_automodel.components.models.mistral_large4.rope_utils.MistralLarge4Rotary.forward(
    positions: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Build differentiable interleaved YaRN tables.

Parameters:

positions

torch.Tensor

Integer tensor of shape [tokens] containing document-relative positions.

Returns: torch.Tensor

FP32 cosine and sine tensors, each of shape [tokens, rotary_dim / 2].

class nemo_automodel.components.models.mistral_large4.rope_utils.SSMax(
    heads: int,
    dtype: torch.dtype
)

Bases: Module

Per-head s * log(position + 1) + b in the query’s compute dtype.

= nn.Parameter(torch.ones(heads, dtype=dtype))

= nn.Parameter(torch.zeros(heads, dtype=dtype))

nemo_automodel.components.models.mistral_large4.rope_utils.SSMax.forward(
    q: torch.Tensor,
    log_positions: torch.Tensor
) -> torch.Tensor

Scale a query without modifying storage shared with another consumer.

Parameters:

torch.Tensor

Tensor of shape [tokens, heads, head_dim].

log_positions

torch.Tensor

FP32 tensor of shape [tokens].

Returns: torch.Tensor

Independent tensor of shape [tokens, heads, head_dim] in q.dtype.

nemo_automodel.components.models.mistral_large4.rope_utils.apply_interleaved_rope(
    x: torch.Tensor,
    tables: tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor

Rotate adjacent channel pairs in FP32, then restore the activation dtype.

Parameters:

torch.Tensor

Tensor of shape [tokens, heads, rotary_dim] or [tokens, rotary_dim].

tables

tuple[torch.Tensor, torch.Tensor]

FP32 cosine and sine tensors of shape [tokens, rotary_dim / 2].

Returns: torch.Tensor

Independent tensor with the same shape and dtype as x.