nemo_automodel.components.models.beam.layers

Beam gated attention, root10 SSMax, rotary precision, and SwiGLU/router overrides.

Module Contents

Classes

Name	Description
`BeamAttention`	Gated attention with per-head q/k RMSNorm and full-width output gating.
`BeamGate`	Normalize top-1 routing like HF, including its epsilon and gradients.
`BeamMLP`	Shared dense MLP with HF’s fp32 SwiGLU boundary, preserving parameter names and MLP dispatch.
`BeamMoE`	Shared MoE with BeamGate installed before loading or parallelization.
`Block`	One Beam decoder layer: sandwich norms + depth-scaled residual adds.

Functions

Name	Description
`_cp_local_positions`	Select the two load-balanced sequence chunks owned by this CP rank.
`apply_rope_hf`	Apply HF Beam’s rotary expression and preserve its backward rounding.
`depth_scaled_residual_coefs`	(sublayer_scale, residual_scale) for the depth-scaled residual add.
`hf_rope_tables`	Expand original FP32 rotary tables into HF’s duplicated-half layout.
`is_moe_layer`	Beam marks every layer `"dense"` or `"sparse"` in `config.mlp_layer_types`.
`layer_rope_parameters`	Return the layer type’s RoPE parameters, or None for NoPE+SSMax.
`packed_sdpa_mask`	Construct a block-diagonal causal mask for packed SDPA.
`reject_flex_attention_mask`	Reject padding on Beam’s unmasked-only flex backend without a device synchronization.
`rotate_half`	Apply HF’s half-vector rotation.
`ssmax_query_scale`	Compute HF Beam’s NoPE query scale in FP32.
`window_size_for_backend`	Translate the visible-token window W: SDPA uses W, TE/FA use inclusive offsets (W-1, 0).

Data

SSMAX_EXPONENT

SSMAX_REFERENCE_LENGTH

API

class nemo_automodel.components.models.beam.layers.BeamAttention(
    config: transformers.models.beam.BeamConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Gated attention with per-head q/k RMSNorm and full-width output gating.

Layer type selects full/sliding attention and RoPE or NoPE+SSMax.

attention_dropout

gate_bias

gate_proj

head_dim

k_norm

k_proj

layer_type

= config.layer_types[layer_idx]

num_heads

= config.num_attention_heads

num_kv_heads

= config.num_key_value_heads

o_proj

q_norm

q_proj

sliding_window

use_rope

use_ssmax

= not self.use_rope

v_proj

window_size

nemo_automodel.components.models.beam.layers.BeamAttention._attend(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attention_mask: torch.Tensor | None,
    qkv_format: str,
    attn_kwargs: dict[str, typing.Any]
) -> torch.Tensor

Run backend attention with the correct packed and sliding-window conventions.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, heads, head_dim] or packed [tokens, heads, head_dim], local under CP and in the compute dtype.

torch.Tensor

Tensor with q’s leading axes and shape […, kv_heads, head_dim].

torch.Tensor

Tensor with the same shape/dtype/device as k.

attention_mask

torch.Tensor | None

Optional [batch, sequence] valid-token mask or backend-supported [batch, 1, sequence, sequence] bool/additive mask.

qkv_format

str

bshd for batched input or thd for packed input.

attn_kwargs

dict[str, Any]

Backend metadata. Integer cu_seqlens/cu_seqlens_padded tensors have shape [documents + 1]; max_seqlen and CP size/rank are scalars.

Returns: torch.Tensor

Attention output with q’s leading axes and either [heads, head_dim]

nemo_automodel.components.models.beam.layers.BeamAttention._sdpa_dropout(
    attn_kwargs: dict[str, typing.Any]
) -> None

Set training-only dropout on the preprocessed SDPA arguments.

Parameters:

attn_kwargs

dict[str, Any]

Mutable SDPA call mapping. Any attn_mask tensor has broadcast layout [batch, heads, query_sequence, key_sequence]; tensor values and storage are unchanged. Only the scalar dropout_p entry is updated.

nemo_automodel.components.models.beam.layers.BeamAttention._set_sliding_cp_transport(
    qkv_format: str
) -> None

Select TE’s CP exchange for this sliding-window layer from the input layout.

Automodel’s CP setup gives sliding-window layers TE’s all_gather exchange. With packed (thd) input that exchange fails before the first step in TE 2.18 (cudaErrorIllegalInstruction), so packed input uses a2a, which matches all_gather numerically; batched (bshd) input keeps all_gather. a2a splits heads across the CP group, so the CP size must divide both head counts. Full-attention layers are not touched.

Parameters:

qkv_format

str

bshd for batched input or thd for packed input.

nemo_automodel.components.models.beam.layers.BeamAttention.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor | None,
    rope_tables: tuple[torch.Tensor, torch.Tensor] | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

Apply gated attention with HF-compatible RoPE or NoPE query scaling.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, hidden] or packed [tokens, hidden], in the compute dtype; sequence/tokens are local under CP.

freqs_cis

torch.Tensor | None

Original FP32 concatenated half-width cosine/sine tensor of shape [batch, sequence, head_dim] or [tokens, head_dim] for unfused RoPE. Fused layout follows position_ids_to_freqs_cis; None for NoPE.

rope_tables

tuple[torch.Tensor, torch.Tensor] | NoneDefaults to None

Optional original FP32 (cos, sin) pair, each of shape [batch, sequence, 1, head_dim] or [tokens, 1, head_dim]. Shared with other layers without mutation. When absent, derived from freqs_cis.

position_ids

torch.Tensor | NoneDefaults to None

Integer tensor of shape [batch, sequence] or [tokens]. Normally local under CP; full-length positions can be selected using cp_size/cp_rank.

attention_mask

torch.Tensor | NoneDefaults to None

Optional [batch, sequence] valid-token mask (1/True = attend), or backend-supported [batch, 1, sequence, sequence] bool/additive mask.

**attn_kwargs

AnyDefaults to {}

Backend metadata: integer cu_seqlens/cu_seqlens_padded tensors of shape [documents + 1], scalar max_seqlen, qkv_format, cp_size/cp_rank.

Returns: torch.Tensor

Tensor with the same shape as x, in the output projection’s compute dtype.

nemo_automodel.components.models.beam.layers.BeamAttention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

Initialize projection and normalization parameters; buffer_device follows the model initializer API.

class nemo_automodel.components.models.beam.layers.BeamGate()

Bases: Gate

Normalize top-1 routing like HF, including its epsilon and gradients.

The override also serves the CUDA-graph routing core via Gate._route_scores.

nemo_automodel.components.models.beam.layers.BeamGate._route_scores(
    scores: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]

Route scores with HF’s normalization for top-1 as well as larger top-k.

Parameters:

scores

torch.Tensor

Tensor of shape [tokens, experts], using the shared Gate score dtype.

Returns: torch.Tensor

Routing weights [tokens, topk], integer expert indices [tokens, topk], and

class nemo_automodel.components.models.beam.layers.BeamMLP(
    dim: int,
    inter_dim: int,
    backend: str,
    dtype: torch.dtype = torch.bfloat16
)

Bases: MLP

Shared dense MLP with HF’s fp32 SwiGLU boundary, preserving parameter names and MLP dispatch.

nemo_automodel.components.models.beam.layers.BeamMLP.forward(
    x: torch.Tensor
) -> torch.Tensor

Compute SwiGLU in FP32 before the down projection.

Parameters:

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions, in the projection compute dtype.

Returns: torch.Tensor

Tensor of shape […, hidden] with the same leading dimensions. The FP32

class nemo_automodel.components.models.beam.layers.BeamMoE(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: MoE

Shared MoE with BeamGate installed before loading or parallelization.

Shared compiled SwiGLU preserves the FP32 activation boundary; Dynamo eager fallback can lose it. Unlike HF, default expert routing weights apply before the down projection and expert outputs accumulate in FP32.

class nemo_automodel.components.models.beam.layers.Block(
    layer_idx: int,
    config: transformers.models.beam.BeamConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

One Beam decoder layer: sandwich norms + depth-scaled residual adds.

attn_scales

input_layernorm

= norm()

layer_type

= config.layer_types[layer_idx]

mlp

= BeamMoE(moe_config, backend)

mlp_scales

post_attention_layernorm

= norm()

post_feedforward_layernorm

= norm()

pre_feedforward_layernorm

= norm()

self_attn

= BeamAttention(config, layer_idx, backend)

nemo_automodel.components.models.beam.layers.Block._mlp(
    x: torch.Tensor,
    padding_mask: torch.Tensor | None
) -> torch.Tensor

Dispatch to the dense MLP or the routed experts.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, hidden] or packed [tokens, hidden], local to this rank under CP, in the layer compute dtype.

padding_mask

torch.Tensor | None

Optional bool tensor matching x’s leading axes; True skips routing.

Returns: torch.Tensor

Tensor with the same shape as x.

nemo_automodel.components.models.beam.layers.Block.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor | None,
    rope_tables: tuple[torch.Tensor, torch.Tensor] | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

Apply attention, feedforward computation and depth-scaled residuals.

Parameters:

torch.Tensor

Tensor [batch, sequence, hidden] or packed [tokens, hidden], local under CP; only x is cast to the FSDP compute dtype.

freqs_cis

torch.Tensor | None

FP32 rotary tensor […, head_dim], aligned with x’s leading axes; concatenated half-width cos/sin when unfused, fused layout otherwise. None for NoPE.

rope_tables

tuple[torch.Tensor, torch.Tensor] | NoneDefaults to None

Optional FP32 (cos, sin), each […, 1, head_dim], shared across layers without mutation or casting; leading axes match x.

position_ids

torch.Tensor | NoneDefaults to None

Integer tensor [batch, sequence] or [tokens], absolute within each document.

attention_mask

torch.Tensor | NoneDefaults to None

Optional [batch, sequence] bool/0-1 mask; True/1 means valid.

padding_mask

torch.Tensor | NoneDefaults to None

Optional bool tensor matching x’s leading axes; True skips routing.

**attn_kwargs

AnyDefaults to {}

Backend metadata: qkv_format (bshd/thd), integer cu_seqlens and cu_seqlens_padded [documents + 1], scalar max_seqlen, cp_size and cp_rank.

Returns: torch.Tensor

Tensor with x’s shape, covering this rank’s local positions.

nemo_automodel.components.models.beam.layers.Block.init_weights(
    buffer_device: torch.device
) -> None

Initialize the layer’s norms, attention, MLP and router buffers on buffer_device.

nemo_automodel.components.models.beam.layers._cp_local_positions(
    position_ids: torch.Tensor,
    local_len: int,
    cp_size: int,
    cp_rank: int
) -> torch.Tensor

Select the two load-balanced sequence chunks owned by this CP rank.

Parameters:

position_ids

torch.Tensor

Integer tensor of shape [batch, global_sequence] or [global_sequence], or the corresponding already-local shape.

local_len

int

Tokens per local sequence, equal to global_sequence / cp_size; must be even when selecting the two head/tail chunks.

cp_size

int

Number of context-parallel ranks.

cp_rank

int

Zero-based context-parallel rank.

Returns: torch.Tensor

Position tensor of shape [batch, local_len] or [local_len]. Already-local

nemo_automodel.components.models.beam.layers.apply_rope_hf(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

Apply HF Beam’s rotary expression and preserve its backward rounding.

Do not eagerly cast x to FP32: HF rounds each product’s gradient to x.dtype.

Parameters:

torch.Tensor

Query/key tensor of shape [batch, sequence, heads, head_dim] or packed [tokens, heads, head_dim], in the activation dtype.

cos

torch.Tensor

Original FP32 cosine tensor of shape [batch, sequence, 1, head_dim] or [tokens, 1, head_dim], broadcastable to x and on the same device.

sin

torch.Tensor

Original FP32 sine tensor with the same shape/device as cos.

Returns: torch.Tensor

New tensor with x’s shape/dtype/device, computed as

nemo_automodel.components.models.beam.layers.depth_scaled_residual_coefs(
    layer_idx: int,
    is_attention: bool
) -> tuple[float, float]

(sublayer_scale, residual_scale) for the depth-scaled residual add.

1-based residual-add index: 2L for attention, 2L + 1 for the MLP; the sublayer output is scaled by 1/sqrt(idx) and the residual stream by sqrt((idx - 1)/idx).

nemo_automodel.components.models.beam.layers.hf_rope_tables(
    freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Expand original FP32 rotary tables into HF’s duplicated-half layout.

Parameters:

freqs_cis

torch.Tensor

FP32 tensor of shape […, head_dim], with arbitrary leading axes corresponding to [batch, sequence] or packed [tokens]. The last axis concatenates half-width cosine and sine values, in that order.

Returns: torch.Tensor

Pair (cos, sin), each a new FP32 tensor of shape […, 1, head_dim] on the

nemo_automodel.components.models.beam.layers.is_moe_layer(
    config: transformers.models.beam.BeamConfig,
    layer_idx: int
) -> bool

Beam marks every layer "dense" or "sparse" in config.mlp_layer_types.

nemo_automodel.components.models.beam.layers.layer_rope_parameters(
    config: transformers.models.beam.BeamConfig,
    layer_type: str
) -> dict[str, typing.Any] | None

Return the layer type’s RoPE parameters, or None for NoPE+SSMax.

Reject flat pre-Beam configs and non-default rope_type rather than silently changing positional encoding.

nemo_automodel.components.models.beam.layers.packed_sdpa_mask(
    cu_seqlens: torch.Tensor,
    total_tokens: int,
    sliding_window: int | None
) -> torch.Tensor

Construct a block-diagonal causal mask for packed SDPA.

Parameters:

cu_seqlens

torch.Tensor

Integer tensor of shape [documents + 1], beginning at zero and ending at total_tokens, with strictly increasing document boundaries.

total_tokens

int

Number of tokens in the concatenated packed buffer.

sliding_window

int | None

Number of visible tokens, including the query, or None.

Returns: torch.Tensor

Bool tensor of shape [1, 1, tokens, tokens] on cu_seqlens.device, where

nemo_automodel.components.models.beam.layers.reject_flex_attention_mask(
    attention_mask: torch.Tensor | None
) -> None

Reject padding on Beam’s unmasked-only flex backend without a device synchronization.

Parameters:

attention_mask

torch.Tensor | None

Optional tensor, normally [batch, sequence]; any non-None value is rejected without inspecting its contents. Use TE or SDPA for padded input.

nemo_automodel.components.models.beam.layers.rotate_half(
    x: torch.Tensor
) -> torch.Tensor

Apply HF’s half-vector rotation.

Parameters:

torch.Tensor

Tensor of shape […, head_dim], with arbitrary leading axes and even head_dim.

Returns: torch.Tensor

New tensor of the same shape/dtype/device, concatenating -x’s second half

nemo_automodel.components.models.beam.layers.ssmax_query_scale(
    position_ids: torch.Tensor
) -> torch.Tensor

Compute HF Beam’s NoPE query scale in FP32.

Parameters:

position_ids

torch.Tensor

Integer tensor of shape [batch, sequence] or packed [tokens], containing absolute per-document positions, local to this rank under CP.

Returns: torch.Tensor

FP32 tensor with the same shape/device as position_ids:

nemo_automodel.components.models.beam.layers.window_size_for_backend(
    sliding_window: int | None,
    attn_impl: str
) -> tuple[int, int] | None

Translate the visible-token window W: SDPA uses W, TE/FA use inclusive offsets (W-1, 0).

Flex takes sliding_window=W directly.

nemo_automodel.components.models.beam.layers.SSMAX_EXPONENT = 0.1

nemo_automodel.components.models.beam.layers.SSMAX_REFERENCE_LENGTH = 2048