nemo_automodel.components.models.praxis.layers

Transformer layers shared by the two towers of the Praxis video backbone.

Every block is pre-norm with a zero-centred RMSNorm (scale 1 + weight), gated grouped-query attention over integer rotary positions, and either a dense SwiGLU or an Automodel MoE feed-forward. The decoder tower modulates each block input with a per-token shift and scale from the encoder.

Norms, rotary rotation, modulation and the router run in float32; everything else in the parameter dtype (bf16 in the released checkpoint).

Module Contents

Classes

Name	Description
`GatedAttention`	Grouped-query attention with a sigmoid output gate per head.
`IntegerRotaryEmbedding`	Rotary (cos, sin) table for large integer positions.
`TransformerBlock`	Pre-norm block `x + attn(norm(x) * s)`, `x + ffn(norm(x) * s)` with `s = 1 / sqrt(layer_index)`.
`TransformerStack`	A stack of `TransformerBlock` sharing one rotary table.
`ZeroCenteredRMSNorm`	RMSNorm whose scale is `1 + weight`; `cast_dtype` selects the compute dtype.

Functions

Name	Description
`apply_rotary`	Rotate `x` with a rotary table in float32 (rotate-half layout).
`build_moe_config`	MoE settings of the Praxis trunk: fp32 sigmoid router with a selection bias, one shared expert.
`modulate`	Return `x * scale + shift` in the dtype of `x`.

Data

ModulationDtype

_compiled_flex_attention

API

class nemo_automodel.components.models.praxis.layers.GatedAttention(
    dim: int,
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
    qk_norm_eps: float,
    attn_backend: str
)

Bases: Module

Grouped-query attention with a sigmoid output gate per head.

Order: project, rotate q and k, RMS-normalize q and k per head, attend, gate, project out.

gate_proj

= nn.Linear(dim, num_heads, bias=False)

k_norm

= ZeroCenteredRMSNorm(head_dim, eps=qk_norm_eps)

k_proj

out_proj

= nn.Linear(num_heads * head_dim, dim, bias=True)

q_norm

= ZeroCenteredRMSNorm(head_dim, eps=qk_norm_eps)

q_proj

= nn.Linear(dim, num_heads * head_dim, bias=False)

v_proj

nemo_automodel.components.models.praxis.layers.GatedAttention.forward(
    x: torch.Tensor,
    rope: torch.Tensor,
    mask: torch.nn.attention.flex_attention.BlockMask | torch.Tensor | None
) -> torch.Tensor

Attend within one sequence.

Parameters:

Tensor

Tensor of shape [batch, tokens, hidden].

rope

Tensor

Float32 rotary table of shape [batch, tokens, head_dim // 2, 2].

mask

BlockMask | Tensor | None

A flex BlockMask over [tokens, tokens], a boolean tensor of shape [tokens, tokens] (true where attention is allowed), or None for full attention.

Returns: Tensor

Tensor of shape [batch, tokens, hidden].

class nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding(
    dim: int,
    effective_dim: int,
    base: int,
    table_size: int
)

Bases: Module

Rotary (cos, sin) table for large integer positions.

Only the first effective_dim // 2 frequency pairs rotate; the remaining pairs have zero frequency. A position p is split as p = high * table_size + low and the angle of each part is formed in float64 so that positions in the millions keep their precision; the two parts are combined with the angle-sum identity.

theta

= self.make_theta(torch.device('cpu'))

nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding.forward(
    positions: torch.Tensor
) -> torch.Tensor

Build the rotary table.

Parameters:

positions

Tensor

Integer tensor of shape [batch, tokens].

Returns: Tensor

Float32 tensor of shape [batch, tokens, dim // 2, 2] holding (cos, sin) per pair.

nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding.make_theta(
    device: torch.device
) -> torch.Tensor

Per-pair frequencies of shape [dim // 2] in float32 on device (zero beyond effective_dim).

class nemo_automodel.components.models.praxis.layers.TransformerBlock(
    dim: int,
    attention: nemo_automodel.components.models.praxis.layers.GatedAttention,
    ffn: torch.nn.Module,
    layer_index: int,
    norm_eps: float,
    modulation_dim: int | None,
    modulation_dtype: nemo_automodel.components.models.praxis.layers.ModulationDtype
)

Bases: Module

Pre-norm block x + attn(norm(x) * s), x + ffn(norm(x) * s) with s = 1 / sqrt(layer_index).

s is a persistent buffer so the checkpoint value (rounded to the parameter dtype) is used.

ln_in

= ZeroCenteredRMSNorm(dim, eps=norm_eps)

ln_out

= ZeroCenteredRMSNorm(dim, eps=norm_eps)

mlp

MLP | MoE

Expose the feed-forward sub-block (ffn) under the name the shared MoE parallelizer looks up.

modulator

Tensor

nemo_automodel.components.models.praxis.layers.TransformerBlock.forward(
    x: torch.Tensor,
    rope: torch.Tensor,
    mask: torch.nn.attention.flex_attention.BlockMask | torch.Tensor | None,
    x_mod: torch.Tensor | None
) -> torch.Tensor

Run the block.

Parameters:

Tensor

Tensor of shape [batch, tokens, hidden].

rope

Tensor

Float32 rotary table of shape [batch, tokens, head_dim // 2, 2].

mask

BlockMask | Tensor | None

Attention mask (see GatedAttention.forward).

x_mod

Tensor | None

Per-token conditioning of shape [batch, tokens, modulation_dim] when the block has a modulator, else None.

Returns: Tensor

Tensor of shape [batch, tokens, hidden].

class nemo_automodel.components.models.praxis.layers.TransformerStack(
    depth: int,
    dense_blocks: tuple[int, ...],
    dim: int,
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
    qk_norm_eps: float,
    norm_eps: float,
    dense_hidden: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    rope: nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding,
    modulation_dim: int | None,
    modulation_dtype: nemo_automodel.components.models.praxis.layers.ModulationDtype,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    dtype: torch.dtype
)

Bases: Module

A stack of TransformerBlock sharing one rotary table.

Blocks listed in dense_blocks use a dense SwiGLU; the others an Automodel MoE.

blocks

= nn.ModuleList(blocks)

nemo_automodel.components.models.praxis.layers.TransformerStack.forward(
    x: torch.Tensor,
    positions: torch.Tensor,
    mask: torch.nn.attention.flex_attention.BlockMask | torch.Tensor | None,
    x_mod: torch.Tensor | None,
    tap_blocks: tuple[int, ...]
) -> tuple[torch.Tensor, list[torch.Tensor]]

Run every block and collect the tap outputs.

Parameters:

Tensor

Tensor of shape [batch, tokens, hidden].

positions

Tensor

Integer tensor of shape [batch, tokens] with the rotary positions.

mask

BlockMask | Tensor | None

Attention mask (see GatedAttention.forward).

x_mod

Tensor | None

Per-token conditioning of shape [batch, tokens, modulation_dim], or None.

tap_blocks

tuple[int, ...]

Indices of the blocks whose outputs are returned.

Returns: Tensor

The output of shape [batch, tokens, hidden] and the tap outputs, each of shape

class nemo_automodel.components.models.praxis.layers.ZeroCenteredRMSNorm(
    dim: int,
    eps: float = 1e-06,
    cast_dtype: torch.dtype | None = torch.float32
)

Bases: Module

RMSNorm whose scale is 1 + weight; cast_dtype selects the compute dtype.

weight

= nn.Parameter(torch.zeros(dim))

nemo_automodel.components.models.praxis.layers.ZeroCenteredRMSNorm.forward(
    x: torch.Tensor
) -> torch.Tensor

Normalize over the last axis.

Parameters:

Tensor

Tensor of shape […, hidden].

Returns: Tensor

Tensor of shape […, hidden] in the dtype of x.

nemo_automodel.components.models.praxis.layers.apply_rotary(
    x: torch.Tensor,
    table: torch.Tensor
) -> torch.Tensor

Rotate x with a rotary table in float32 (rotate-half layout).

Parameters:

Tensor

Tensor of shape [batch, heads, tokens, head_dim]; channel j pairs with j + head_dim // 2.

table

Tensor

Float32 tensor of shape [batch, tokens, head_dim // 2, 2] from IntegerRotaryEmbedding.

Returns: Tensor

Tensor of shape [batch, heads, tokens, head_dim] in the dtype of x.

nemo_automodel.components.models.praxis.layers.build_moe_config(
    dim: int,
    expert_hidden: int,
    num_routed_experts: int,
    top_k_routed: int,
    routed_scaling_factor: float,
    dtype: torch.dtype,
    gate_bias_update_factor: float,
    aux_loss_coeff: float = 0.0
) -> nemo_automodel.components.moe.config.MoEConfig

MoE settings of the Praxis trunk: fp32 sigmoid router with a selection bias, one shared expert.

The routed sum is scaled by routed_scaling_factor and the top-k scores are renormalized to sum to one, as in DeepSeek-V3.

nemo_automodel.components.models.praxis.layers.modulate(
    x: torch.Tensor,
    shift: torch.Tensor,
    scale: torch.Tensor,
    dtype: nemo_automodel.components.models.praxis.layers.ModulationDtype
) -> torch.Tensor

Return x * scale + shift in the dtype of x.

Parameters:

Tensor

Tensor of shape [batch, tokens, hidden].

shift

Tensor

Tensor broadcastable to x.

scale

Tensor

Tensor broadcastable to x.

dtype

ModulationDtype

"mixed" and "float32" compute in float32, "bfloat16" in bf16.

Returns: Tensor

Tensor of shape [batch, tokens, hidden] in the dtype of x.

nemo_automodel.components.models.praxis.layers.ModulationDtype = Literal['bfloat16', 'float32', 'mixed']

nemo_automodel.components.models.praxis.layers._compiled_flex_attention = torch.compile(flex_attention)