nemo_automodel.components.models.beam.model

Native Beam decoder: gated attention, RoPE/NoPE with root10 SSMax, sandwich norms, depth-scaled residuals, dense/MoE SwiGLU, and soft-capped logits. The module tree preserves HF names; BeamStateDictAdapter converts the grouped expert weights.

Module Contents

Classes

Name	Description
`BeamForCausalLM`	Beam causal LM on Automodel’s MoE stack (see module docstring).
`BeamModel`	Embedding, decoder layers and final norm (the `model.` subtree of the HF checkpoint).

Functions

Name	Description
`beam_moe_config`	Map Beam’s sigmoid/grouped-top-k router and SwiGLU experts to MoEConfig.
`resolve_backend_gate_precision`	Default router compute to fp32 and reject lower precision to preserve HF expert selection.

Data

FP32_SKIP_MODULES

ModelClass

API

class nemo_automodel.components.models.beam.model.BeamForCausalLM(
    config: transformers.models.beam.BeamConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Beam causal LM on Automodel’s MoE stack (see module docstring).

_keep_in_fp32_modules_strict

list[str] = ['mlp.gate.weight', 'e_score_correction_bias']

_pp_keep_self_forward

bool = True

backend

final_logit_softcapping

= getattr(config, 'final_logit_softcapping', None)

lm_head

model

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

nemo_automodel.components.models.beam.model.BeamForCausalLM.forward(
    input_ids: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast

Return soft-capped logits or pipeline-stage hidden states.

Parameters:

input_ids

torch.Tensor

Integer tensor [batch, sequence] or packed [tokens]; later PP stages receive floating hidden states with an extra final hidden axis. Positions are local under CP.

position_ids

torch.Tensor | NoneDefaults to None

Optional 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. None on flex.

padding_mask

torch.Tensor | NoneDefaults to None

Optional bool tensor matching the position axes; True skips routing.

logits_to_keep

int | torch.TensorDefaults to 0

Zero for all positions, a positive final-position count, or integer tensor [selected_positions] indexing the sequence axis.

output_hidden_states

bool | NoneDefaults to None

Whether to return the final hidden states.

**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: CausalLMOutputWithPast

logits [batch, selected_sequence, vocab] or packed

nemo_automodel.components.models.beam.model.BeamForCausalLM.from_config(
    config: transformers.models.beam.BeamConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.beam.model.BeamForCausalLM

classmethod

Construct the native model from an HF BeamConfig and optional backend/MoE overrides.

nemo_automodel.components.models.beam.model.BeamForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.beam.model.BeamForCausalLM

classmethod

Construct from the checkpoint’s config; framework checkpoint loading supplies the weights.

nemo_automodel.components.models.beam.model.BeamForCausalLM.get_input_embeddings() -> torch.nn.Module | None

Return this stage’s embedding module, or None when another stage owns it.

nemo_automodel.components.models.beam.model.BeamForCausalLM.get_output_embeddings() -> torch.nn.Module | None

Return this stage’s LM head, or None when another stage owns it.

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

Initialize weights and cast to dtype while preserving original FP32 router values.

nemo_automodel.components.models.beam.model.BeamForCausalLM.set_input_embeddings(
    value: torch.nn.Module | None
) -> None

Replace this stage’s embedding module; None marks a later pipeline stage.

nemo_automodel.components.models.beam.model.BeamForCausalLM.set_output_embeddings(
    new_embeddings: torch.nn.Module | None
) -> None

Replace this stage’s LM head; None marks a non-final pipeline stage.

class nemo_automodel.components.models.beam.model.BeamModel(
    config: transformers.models.beam.BeamConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict[str, typing.Any] | None = None
)

Bases: Module

Embedding, decoder layers and final norm (the model. subtree of the HF checkpoint).

embed_tokens

head_dim

layers

= torch.nn.ModuleDict()

max_seq_len

= config.max_position_embeddings

moe_config

norm

rotary_emb

= nn.ModuleDict()

nemo_automodel.components.models.beam.model.BeamModel.forward(
    input_ids: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

Run this pipeline stage with one shared FP32 RoPE-table pair per layer type.

Parameters:

input_ids

torch.Tensor

Integer tensor [batch, sequence] or packed [tokens]; later PP stages receive floating hidden states with an extra final hidden axis. Positions are local under CP.

position_ids

torch.Tensor | NoneDefaults to None

Optional 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. None on flex.

padding_mask

torch.Tensor | NoneDefaults to None

Optional bool tensor matching the position 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

Hidden states [batch, sequence, hidden] or [tokens, hidden], local under CP,

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

Initialize this stage’s parameters and select the rotary-buffer device.

nemo_automodel.components.models.beam.model.beam_moe_config(
    config: transformers.models.beam.BeamConfig,
    model_dtype: torch.dtype,
    overrides: dict[str, typing.Any] | None = None
) -> nemo_automodel.components.moe.config.MoEConfig

Map Beam’s sigmoid/grouped-top-k router and SwiGLU experts to MoEConfig.

Correction bias affects selection only and is never updated. Router storage and compute remain fp32. Automodel applies routed_scaling_factor to routing weights; HF applies it after expert aggregation, differing only in rounding.

nemo_automodel.components.models.beam.model.resolve_backend_gate_precision(
    backend: nemo_automodel.components.models.common.BackendConfig
) -> nemo_automodel.components.models.common.BackendConfig

Default router compute to fp32 and reject lower precision to preserve HF expert selection.

nemo_automodel.components.models.beam.model.FP32_SKIP_MODULES = ('gate',)

nemo_automodel.components.models.beam.model.ModelClass = BeamForCausalLM