bridge.models.beam.beam_config

Beam model configuration: a HybridModel config plus the Beam architecture fields.

Module Contents

Classes

BeamTransformerConfig	Bridge TransformerConfig plus the Beam fields Megatron-Core does not declare.
BeamModelConfig	HybridModel config whose builder uses the Beam stack spec and soft-caps the output layer.

Functions

beam_pipeline_layer_pattern

Split a Beam hybrid pattern into num_stages |-separated segments.

API

class bridge.models.beam.beam_config.BeamTransformerConfig

Bases: megatron.bridge.models.transformer_config.TransformerConfig

Bridge TransformerConfig plus the Beam fields Megatron-Core does not declare.

The modules in :mod:~megatron.bridge.models.beam.modeling_beam read them.

attention_output_gate_bias: bool

False

Gate the attention output with sigmoid(gate + bias[head]); requires attention_output_gate. The bias lives in SelfAttention.output_gate_bias and is sharded across TP with the heads.

post_self_attn_layernorm: bool

False

Normalize the self-attention output before the residual add (“sandwich” norm).

post_mlp_layernorm: bool

False

Normalize the MLP output before the residual add (“sandwich” norm).

depth_scaled_residual: bool

False

Replace each residual add with out*(1/sqrt(idx)) + residual*sqrt((idx-1)/idx), idx = layer_number + 1 with the hybrid stack’s 1-based layer_number.

ssmax: bool

False

Scale queries on NoPE layers by (1 + (position + 1) / 2048) ** 0.1; requires no_rope_freq.

moe_router_fp32_params: bool

False

Keep router weights and biases in fp32, including through Float16Module; requires moe_router_dtype='fp32'.

final_logit_softcapping: Optional[float]

None

logits = cap * tanh(logits / cap) on the output layer; None or 0 disables it.

finalize() → None

Run the deferred Megatron-Core post-init, then validate the Beam fields.

bridge.models.beam.beam_config.beam_pipeline_layer_pattern(pattern: str, num_stages: int) → str

Split a Beam hybrid pattern into num_stages |-separated segments.

Beam layer i is the hybrid pair pattern[2i:2i+2] (attention, then MLP/MoE); pairs stay on one stage and earlier stages take the remainder. Existing separators are dropped, so a pattern saved under another pipeline size is re-split.

class bridge.models.beam.beam_config.BeamModelConfig

Bases: megatron.bridge.models.common.ModelConfigOverrideMixin, megatron.bridge.models.hybrid.hybrid_builder.HybridModelConfig

HybridModel config whose builder uses the Beam stack spec and soft-caps the output layer.

hybrid_layer_pattern holds two hybrid layers per Beam layer: * attention followed by - (dense MLP) or E (MoE), so transformer.num_layers and the per-layer lists (no_rope_freq, window_attn_skip_freq) count hybrid layers. _builder_ names

Class:

~.beam_builder.BeamModelBuilder.

builder: ClassVar[str]

‘megatron.bridge.models.beam.BeamModelBuilder’

transformer_config_class: ClassVar[type[megatron.bridge.models.transformer_config.TransformerConfig]]

None

finalize() → None

Split the pattern at Beam-layer boundaries for pipeline parallelism, then validate.

get_builder_cls() → type

Resolve the Beam builder through Bridge’s target allowlist.

as_dict() → dict[str, Any]

Serialize the config with a symbolic activation function.

classmethod from_dict(

data: dict[str, Any],

) → bridge.models.beam.beam_config.BeamModelConfig

Deserialize a Beam config and restore its activation callable.