nemo_automodel.components.models.laguna.model

Module Contents

Classes

Name	Description
`LagunaAttention`	Laguna attention: explicit per-layer heads, QK RMSNorm, and output gating.
`LagunaBlock`	Decoder block that uses dense MLP for configured layers and MoE otherwise.
`LagunaForCausalLM`	Causal LM wrapper for Laguna with Automodel checkpoint adapters.
`LagunaModel`	Backbone model for Laguna SFT.
`LagunaRMSNorm`	RMSNorm with fp32 variance, matching the Laguna reference implementation.
`LagunaRotaryEmbedding`	Rotary embedding with Laguna’s nested full-attention/SWA RoPE support.

Functions

Name	Description
`_apply_rotary_pos_emb`	Apply RoPE to query and key states.
`_config_with_rope`	-
`_convert_bool_4d_mask_to_additive`	Convert a bool 4D attention mask to additive attention-mask layout.
`_derive_padding_mask`	Derive an MoE padding mask from a sequence or attention mask.
`_eager_attention_forward`	Run eager attention over projected query/key/value states.
`_ensure_additive_mask`	Normalize an optional precomputed mask into additive layout.
`_fallback_additive_mask`	Build a causal additive attention mask.
`_normalize_gating_mode`	-
`_repeat_kv`	Repeat key/value heads for grouped-query attention.
`_rotate_half`	Rotate the two halves of the RoPE feature axis.
`_sdpa_attention_forward`	Run SDPA attention over projected query/key/value states.

Data

ModelClass

__all__

API

class nemo_automodel.components.models.laguna.model.LagunaAttention(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    layer_idx: int
)

Bases: Module

Laguna attention: explicit per-layer heads, QK RMSNorm, and output gating.

attention_dropout

= float(config.attention_dropout or 0.0)

attention_type

g_proj

gating_mode

= _normalize_gating_mode(gating)

head_dim

= config.head_dim

is_sliding

= self.attention_type == 'sliding_attention'

k_norm

k_proj

num_heads

num_key_value_groups

= self.num_heads // self.num_key_value_heads

num_key_value_heads

= config.num_key_value_heads

o_proj

q_norm

q_proj

scaling

= self.head_dim ** -0.5

sliding_window

= config.sliding_window if self.is_sliding else None

v_proj

nemo_automodel.components.models.laguna.model.LagunaAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]

Run Laguna attention for one decoder layer.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

position_embeddings

tuple[torch.Tensor, torch.Tensor]

Tuple of cosine and sine RoPE tensors, each of shape [batch, sequence, rotary_dim].

attention_mask

torch.Tensor | NoneDefaults to None

Optional additive mask broadcastable to [batch, heads, sequence, key_sequence].

**kwargs

AnyDefaults to {}

Additional attention backend arguments.

Returns: torch.Tensor

Tuple of attention output [batch, sequence, hidden] and optional attention weights

nemo_automodel.components.models.laguna.model.LagunaAttention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

class nemo_automodel.components.models.laguna.model.LagunaBlock(
    layer_idx: int,
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Decoder block that uses dense MLP for configured layers and MoE otherwise.

attention_type

= config.layer_types[layer_idx]

input_layernorm

mlp

= MoE(moe_config, backend)

post_attention_layernorm

self_attn

nemo_automodel.components.models.laguna.model.LagunaBlock.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Run a decoder block.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

attention_mask

torch.Tensor | NoneDefaults to None

Optional additive attention mask broadcastable to [batch, heads, sequence, key_sequence].

position_embeddings

tuple[torch.Tensor, torch.Tensor]

Tuple of cosine and sine RoPE tensors, each of shape [batch, sequence, rotary_dim].

padding_mask

torch.Tensor | NoneDefaults to None

Optional bool tensor of shape [batch, sequence], where True marks padding.

**kwargs

AnyDefaults to {}

Additional attention backend arguments.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden].

nemo_automodel.components.models.laguna.model.LagunaBlock.init_weights(
    buffer_device: torch.device
) -> None

class nemo_automodel.components.models.laguna.model.LagunaForCausalLM(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Causal LM wrapper for Laguna with Automodel checkpoint adapters.

_keep_in_fp32_modules_strict

= ['mlp.gate.e_score_correction_bias', 'rotary_emb']

backend

= backend or BackendConfig()

lm_head

model

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    past_key_values: typing.Any = None,
    use_cache: bool | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: typing.Optional[bool] = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast

Run the Laguna causal language model.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Optional token IDs of shape [batch, sequence].

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional embeddings of shape [batch, sequence, hidden].

position_ids

torch.Tensor | NoneDefaults to None

Optional position IDs of shape [batch, sequence].

attention_mask

torch.Tensor | dict[str, torch.Tensor] | NoneDefaults to None

Optional 2D bool/int mask of shape [batch, sequence], a 4D additive mask, or a mapping with per-attention-type masks keyed by “full_attention” and “sliding_attention”.

padding_mask

torch.Tensor | NoneDefaults to None

Optional bool tensor of shape [batch, sequence], where True marks tokens excluded from MoE routing.

past_key_values

AnyDefaults to None

Unsupported KV-cache state.

use_cache

bool | NoneDefaults to None

Unsupported cache flag.

logits_to_keep

Union[int, torch.Tensor]Defaults to 0

If 0, compute logits for all sequence positions. If an int or tensor, compute logits only for the selected trailing positions.

output_hidden_states

Optional[bool]Defaults to None

When true, include final hidden states in the output.

**kwargs

AnyDefaults to {}

Additional attention backend arguments.

Returns: CausalLMOutputWithPast

Causal LM output with logits and optional hidden states.

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.from_config(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

classmethod

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)

classmethod

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.get_input_embeddings()

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.get_output_embeddings()

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.set_input_embeddings(
    value
)

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.set_output_embeddings(
    new_embeddings
)

nemo_automodel.components.models.laguna.model.LagunaForCausalLM.update_moe_gate_bias() -> None

class nemo_automodel.components.models.laguna.model.LagunaModel(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict | None = None
)

Bases: Module

Backbone model for Laguna SFT.

embed_tokens

has_sliding_layers

= 'sliding_attention' in config.layer_types

layers

moe_config

= moe_config or MoEConfig(**moe_defaults)

norm

rotary_emb

swa_rotary_emb

nemo_automodel.components.models.laguna.model.LagunaModel._build_causal_mask_mapping(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None,
    position_ids: torch.Tensor
) -> dict[str, torch.Tensor]

Build additive attention masks for full and sliding Laguna layers.

Parameters:

inputs_embeds

torch.Tensor

Input embedding tensor of shape [batch, sequence, hidden].

attention_mask

torch.Tensor | dict[str, torch.Tensor] | None

Optional sequence mask of shape [batch, sequence], 4D bool/additive mask, or mapping with masks keyed by “full_attention” and “sliding_attention”.

position_ids

torch.Tensor

Position tensor of shape [batch, sequence].

Returns: dict[str, torch.Tensor]

Mapping from attention type to additive masks of shape

nemo_automodel.components.models.laguna.model.LagunaModel.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Run the Laguna decoder stack.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Optional token IDs of shape [batch, sequence].

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional embeddings of shape [batch, sequence, hidden].

position_ids

torch.Tensor | NoneDefaults to None

Optional position IDs of shape [batch, sequence].

attention_mask

torch.Tensor | dict[str, torch.Tensor] | NoneDefaults to None

Optional 2D bool/int mask of shape [batch, sequence], a 4D additive mask, or a mapping with per-attention-type masks keyed by “full_attention” and “sliding_attention”.

padding_mask

torch.Tensor | NoneDefaults to None

Optional bool tensor of shape [batch, sequence], where True marks tokens excluded from MoE routing.

**kwargs

AnyDefaults to {}

Additional attention backend arguments.

Returns: torch.Tensor

Final hidden states of shape [batch, sequence, hidden].

nemo_automodel.components.models.laguna.model.LagunaModel.init_weights(
    buffer_device: torch.device | None = None
) -> None

class nemo_automodel.components.models.laguna.model.LagunaRMSNorm(
    hidden_size: int,
    eps: float = 1e-06,
    dtype: torch.dtype = torch.bfloat16
)

Bases: Module

RMSNorm with fp32 variance, matching the Laguna reference implementation.

weight

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

nemo_automodel.components.models.laguna.model.LagunaRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Apply RMS normalization while computing variance in fp32.

Parameters:

hidden_states

torch.Tensor

Tensor of shape […, hidden].

Returns: torch.Tensor

Tensor of shape […, hidden].

nemo_automodel.components.models.laguna.model.LagunaRMSNorm.reset_parameters() -> None

class nemo_automodel.components.models.laguna.model.LagunaRotaryEmbedding(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    device: torch.device | None = None
)

Bases: Module

Rotary embedding with Laguna’s nested full-attention/SWA RoPE support.

inv_freq

Tensor

max_seq_len_cached

= config.max_position_embeddings

original_max_seq_len

= config.max_position_embeddings

rope_type

= config.rope_parameters['rope_type']

nemo_automodel.components.models.laguna.model.LagunaRotaryEmbedding._compute_default_rope_parameters(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    device: torch.device | None = None,
    seq_len: int | None = None
) -> tuple[torch.Tensor, float]

staticmethod

nemo_automodel.components.models.laguna.model.LagunaRotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Compute RoPE cosine and sine tables.

Parameters:

torch.Tensor

Activation tensor of shape [batch, sequence, hidden], used for device and dtype.

position_ids

torch.Tensor

Position tensor of shape [batch, sequence].

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

Tuple of cosine and sine tensors, each of shape [batch, sequence, rotary_dim].

nemo_automodel.components.models.laguna.model._apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Apply RoPE to query and key states.

Parameters:

torch.Tensor

Query tensor of shape [batch, heads, sequence, head_dim].

torch.Tensor

Key tensor of shape [batch, key_value_heads, sequence, head_dim].

cos

torch.Tensor

Cosine tensor of shape [batch, sequence, rotary_dim].

sin

torch.Tensor

Sine tensor of shape [batch, sequence, rotary_dim].

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

Tuple of rotated query and key tensors with the same shapes as q and k.

nemo_automodel.components.models.laguna.model._config_with_rope(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    rope_parameters: dict[str, typing.Any]
) -> nemo_automodel.components.models.laguna.config.LagunaConfig

nemo_automodel.components.models.laguna.model._convert_bool_4d_mask_to_additive(
    attention_mask: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor

Convert a bool 4D attention mask to additive attention-mask layout.

Parameters:

attention_mask

torch.Tensor

Tensor of shape [batch, heads_or_one, sequence, key_sequence]. Bool masks use True for visible entries; non-bool masks are already additive.

dtype

torch.dtype

Floating-point dtype for the additive mask.

Returns: torch.Tensor

Additive mask tensor with the same shape, or the original non-bool mask.

nemo_automodel.components.models.laguna.model._derive_padding_mask(
    attention_mask: torch.Tensor
) -> torch.Tensor

Derive an MoE padding mask from a sequence or attention mask.

Parameters:

attention_mask

torch.Tensor

Tensor of shape [batch, sequence] with nonzero valid tokens, or [batch, heads_or_one, sequence, key_sequence] in bool or additive mask layout.

Returns: torch.Tensor

Bool tensor of shape [batch, sequence], where True marks padding.

nemo_automodel.components.models.laguna.model._eager_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor]

Run eager attention over projected query/key/value states.

Parameters:

module

nn.Module

Attention module carrying head-repeat metadata.

query

torch.Tensor

Query tensor of shape [batch, heads, sequence, head_dim].

key

torch.Tensor

Key tensor of shape [batch, key_value_heads, key_sequence, head_dim].

value

torch.Tensor

Value tensor of shape [batch, key_value_heads, key_sequence, head_dim].

attention_mask

torch.Tensor | None

Optional additive mask broadcastable to [batch, heads, sequence, key_sequence].

scaling

float

Multiplicative scale for attention logits.

dropout

floatDefaults to 0.0

Attention dropout probability.

**kwargs

AnyDefaults to {}

Ignored attention backend arguments.

Returns: torch.Tensor

Tuple of context tensor [batch, sequence, heads, head_dim] and attention weights

nemo_automodel.components.models.laguna.model._ensure_additive_mask(
    mask: torch.Tensor | None,
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None,
    sliding_window: int | None
) -> torch.Tensor

Normalize an optional precomputed mask into additive layout.

Parameters:

mask

torch.Tensor | None

Optional attention mask of shape [batch, heads_or_one, sequence, key_sequence]. Bool masks use True for visible entries; additive masks are passed through.

batch_size

int

Batch size used when mask is absent.

seq_len

int

Query and key sequence length used when mask is absent.

dtype

torch.dtype

Floating-point dtype for the additive mask.

device

torch.device

Device for the additive mask.

attention_mask

torch.Tensor | None

Optional sequence mask of shape [batch, sequence] with nonzero valid tokens.

sliding_window

int | None

Optional sliding-window size for local attention.

Returns: torch.Tensor

Additive mask tensor of shape [batch, heads_or_one, sequence, key_sequence].

nemo_automodel.components.models.laguna.model._fallback_additive_mask(
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor

Build a causal additive attention mask.

Parameters:

batch_size

int

Batch size for the returned mask.

seq_len

int

Query and key sequence length.

dtype

torch.dtype

Floating-point dtype for the additive mask.

device

torch.device

Device for the returned mask.

attention_mask

torch.Tensor | NoneDefaults to None

Optional sequence mask of shape [batch, sequence] with nonzero valid tokens.

sliding_window

int | NoneDefaults to None

Optional sliding-window size for local attention.

Returns: torch.Tensor

Additive mask tensor of shape [batch, 1, sequence, sequence].

nemo_automodel.components.models.laguna.model._normalize_gating_mode(
    gating: bool | str
) -> str | None

nemo_automodel.components.models.laguna.model._repeat_kv(
    hidden_states: torch.Tensor,
    n_rep: int
) -> torch.Tensor

Repeat key/value heads for grouped-query attention.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, key_value_heads, sequence, head_dim].

n_rep

int

Number of repeats per key/value head.

Returns: torch.Tensor

Tensor of shape [batch, key_value_heads * n_rep, sequence, head_dim].

nemo_automodel.components.models.laguna.model._rotate_half(
    x: torch.Tensor
) -> torch.Tensor

Rotate the two halves of the RoPE feature axis.

Parameters:

torch.Tensor

Tensor of shape […, rotary_dim], where rotary_dim is even.

Returns: torch.Tensor

Tensor of shape […, rotary_dim].

nemo_automodel.components.models.laguna.model._sdpa_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, None]

Run SDPA attention over projected query/key/value states.

Parameters:

module

nn.Module

Attention module carrying head-repeat metadata.

query

torch.Tensor

Query tensor of shape [batch, heads, sequence, head_dim].

key

torch.Tensor

Key tensor of shape [batch, key_value_heads, key_sequence, head_dim].

value

torch.Tensor

Value tensor of shape [batch, key_value_heads, key_sequence, head_dim].

attention_mask

torch.Tensor | None

Optional additive mask broadcastable to [batch, heads, sequence, key_sequence].

scaling

float

Ignored because SDPA applies its own scale.

dropout

floatDefaults to 0.0

Attention dropout probability.

**kwargs

AnyDefaults to {}

Ignored attention backend arguments.

Returns: tuple[torch.Tensor, None]

Tuple of context tensor [batch, sequence, heads, head_dim] and None for weights.

nemo_automodel.components.models.laguna.model.ModelClass = LagunaForCausalLM

nemo_automodel.components.models.laguna.model.__all__ = ['LagunaForCausalLM', 'LagunaModel', 'ModelClass']