Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsnemo_automodel.components.models.beam

nemo_automodel.components.models.beam

View as Markdown

Submodules

Package Contents

Classes

NameDescription
BeamForCausalLMBeam causal LM on Automodel’s MoE stack (see module docstring).
BeamStateDictAdapterConverts between Hugging Face Beam checkpoints and the grouped-experts native format.

Data

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

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.state_dict_adapter.BeamStateDictAdapter()

Bases: Qwen3MoeStateDictAdapter

Converts between Hugging Face Beam checkpoints and the grouped-experts native format.

Beam stores its experts exactly like Qwen3-MoE, so the Qwen3-MoE conversion (including fused PEFT v5 ParamWrapper LoRA export/import and low-memory DCP loading) applies unchanged. The native module tree mirrors the HF one, so every tensor outside the experts keeps its HF name and passes through unchanged (attention projections and gate_bias, q/k norms, the four per-layer norms, the dense MLP, mlp.gate.weight and the mlp.gate.e_score_correction_bias buffer, embeddings, final norm, lm_head). Experts differ:

HF checkpoint (per expert, the qwen2_moe conversion mapping)::

model.layers.{L}.mlp.experts.{E}.gate_proj.weight [moe_inter_dim, hidden] model.layers.{L}.mlp.experts.{E}.up_proj.weight [moe_inter_dim, hidden] model.layers.{L}.mlp.experts.{E}.down_proj.weight [hidden, moe_inter_dim]

Native (grouped)::

model.layers.{L}.mlp.experts.gate_and_up_projs [n_experts, hidden, 2 * moe_inter_dim] model.layers.{L}.mlp.experts.down_projs [n_experts, moe_inter_dim, hidden]

to_hf therefore writes the per-expert layout that transformers itself produces with save_pretrained, not the stacked in-memory layout.

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