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

## Module Contents

### Classes

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

### Functions

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

### Data

[`ModelClass`](#nemo_automodel-components-models-laguna-model-ModelClass)

[`__all__`](#nemo_automodel-components-models-laguna-model-__all__)

### API

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

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

Tensor of shape \[batch, sequence, hidden].

Tuple of cosine and sine RoPE tensors, each of shape
\[batch, sequence, rotary\_dim].

Optional additive mask broadcastable to
\[batch, heads, sequence, key\_sequence].

Additional attention backend arguments.

**Returns:** `torch.Tensor`

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

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

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

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

Tensor of shape \[batch, sequence, hidden].

Optional additive attention mask broadcastable to
\[batch, heads, sequence, key\_sequence].

Tuple of cosine and sine RoPE tensors, each of shape
\[batch, sequence, rotary\_dim].

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

Additional attention backend arguments.

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

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

```python
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](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

Causal LM wrapper for Laguna with Automodel checkpoint adapters.

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

Optional token IDs of shape \[batch, sequence].

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

Optional position IDs of shape \[batch, sequence].

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

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

Unsupported KV-cache state.

Unsupported cache flag.

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

When true, include final hidden states in the output.

Additional attention backend arguments.

**Returns:** `CausalLMOutputWithPast`

Causal LM output with logits and optional hidden states.

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

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

classmethod

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

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

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

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

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

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

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

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

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

Optional sequence mask of shape \[batch, sequence], 4D bool/additive
mask, or mapping with masks keyed by "full\_attention" and "sliding\_attention".

Position tensor of shape \[batch, sequence].

**Returns:** `dict[str, torch.Tensor]`

Mapping from attention type to additive masks of shape

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

Optional token IDs of shape \[batch, sequence].

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

Optional position IDs of shape \[batch, sequence].

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

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

Additional attention backend arguments.

**Returns:** `torch.Tensor`

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

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

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

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

Apply RMS normalization while computing variance in fp32.

**Parameters:**

Tensor of shape \[..., hidden].

**Returns:** `torch.Tensor`

Tensor of shape \[..., hidden].

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

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

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

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

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

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

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

Query tensor of shape \[batch, heads, sequence, head\_dim].

Key tensor of shape \[batch, key\_value\_heads, sequence, head\_dim].

Cosine tensor of shape \[batch, sequence, rotary\_dim].

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

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

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

Tensor of shape \[batch, heads\_or\_one, sequence, key\_sequence].
Bool masks use True for visible entries; non-bool masks are already additive.

Floating-point dtype for the additive mask.

**Returns:** `torch.Tensor`

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

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

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.

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

Attention module carrying head-repeat metadata.

Query tensor of shape \[batch, heads, sequence, head\_dim].

Key tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

Value tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

Optional additive mask broadcastable to
\[batch, heads, sequence, key\_sequence].

Multiplicative scale for attention logits.

Attention dropout probability.

Ignored attention backend arguments.

**Returns:** `torch.Tensor`

Tuple of context tensor \[batch, sequence, heads, head\_dim] and attention weights

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

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 used when `mask` is absent.

Query and key sequence length used when `mask` is absent.

Floating-point dtype for the additive mask.

Device for the additive mask.

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

Optional sliding-window size for local attention.

**Returns:** `torch.Tensor`

Additive mask tensor of shape \[batch, heads\_or\_one, sequence, key\_sequence].

```python
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 for the returned mask.

Query and key sequence length.

Floating-point dtype for the additive mask.

Device for the returned mask.

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

Optional sliding-window size for local attention.

**Returns:** `torch.Tensor`

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

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

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

Tensor of shape \[batch, key\_value\_heads, sequence, head\_dim].

Number of repeats per key/value head.

**Returns:** `torch.Tensor`

Tensor of shape \[batch, key\_value\_heads \* n\_rep, sequence, head\_dim].

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

Rotate the two halves of the RoPE feature axis.

**Parameters:**

Tensor of shape \[..., rotary\_dim], where rotary\_dim is even.

**Returns:** `torch.Tensor`

Tensor of shape \[..., rotary\_dim].

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

Attention module carrying head-repeat metadata.

Query tensor of shape \[batch, heads, sequence, head\_dim].

Key tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

Value tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

Optional additive mask broadcastable to
\[batch, heads, sequence, key\_sequence].

Ignored because SDPA applies its own scale.

Attention dropout probability.

Ignored attention backend arguments.

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

Tuple of context tensor \[batch, sequence, heads, head\_dim] and `None` for weights.

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

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