`bridge.models.beam.modeling_beam`#

Beam modules on top of an unmodified Megatron-Core HybridModel.

Each Beam decoder layer becomes two HybridStack layers: an attention-only layer (*) followed by a dense MLP (-) or MoE (E) layer. Megatron-Core’s layers cover Beam except for:

SSMax query scaling on NoPE layers,
a learnable per-head bias on the attention output gate,
sandwich norms (a norm on each sublayer output before the residual add),
depth-scaled residual adds,
fp32 router parameters (and two router-GEMM fixes they need),
fp32 final-logit soft-capping.

Everything except the last item is swapped into the stack spec by

func:: beam_hybrid_stack_spec, which builds the hybrid layers from Megatron-Core’s Transformer Engine GPT layer submodules with the subclasses below. Megatron-Core offers no hook for one piece: the distributed optimizer’s parameter-group index map, which assumes a group holds a single parameter dtype.
func:: patch_distributed_optimizer_group_index_map monkey patches it once the first fp32 Beam router is built.

Module Contents#

Classes#

`BeamOutputLayer`	`ColumnParallelLinear` whose output is soft-capped in fp32.
`_BeamSSMaxCoreAttentionMixin`	Core attention that applies SSMax to the query of NoPE layers.
`BeamTEDotProductAttention`	Transformer Engine `TEDotProductAttention` with SSMax on NoPE layers.
`AttentionOutputGateBias`	Zero-initialized, TP-sharded per-head bias for the attention output gate.
`BeamSelfAttention`	`SelfAttention` whose output gate carries a per-head bias.
`BeamTransformerLayerSubmodules`	`TransformerLayerSubmodules` plus the two sandwich norms.
`BeamTransformerLayerMixin`	Sandwich norms and depth-scaled residual adds for an attention-only or MLP-only hybrid layer.
`BeamTransformerLayer`	Attention-only hybrid layer (`*`) with Beam’s sandwich norm and depth-scaled residual.
`BeamMLPLayer`	Dense MLP hybrid layer (`-`) with Beam’s sandwich norm and depth-scaled residual.
`BeamMoETransformerLayer`	MoE hybrid layer (`E`) with Beam’s sandwich norm and depth-scaled residual.
`_BeamRouterGatingLinear`	`RouterGatingLinearFunction` for routers whose weight dtype differs from the input dtype.
`BeamTopKRouter`	`TopKRouter` that keeps its weight (and bias) in fp32 when `moe_router_fp32_params` is set.

Functions#

`beam_logit_softcapping`	`cap * tanh(logits / cap)` computed and returned in fp32, exactly as `BeamForCausalLM` does.
`validate_beam_config`	Reject Beam features on code paths that would silently skip them.
`ssmax_absolute_positions`	Absolute sequence positions for the tokens this rank actually holds.
`ssmax_packed_positions`	Within-sequence positions for every token of a packed (`thd`) buffer.
`apply_ssmax_to_query`	Scale `query` by `(1 + (position + 1) / 2048) ** 0.1` for each of its positions.
`_is_ssmax_layer`	SSMax applies to the NoPE layers (`no_rope_freq` is 0-indexed, `layer_number` 1-indexed).
`apply_output_gate`	Gate `x` with `sigmoid(gate [+ bias per head])`.
`depth_scaled_residual_coefs`	`(sublayer_scale, residual_scale)` for the residual add of one hybrid layer; the squares sum to one.
`scaled_residual_add`	`output * coefs[0] + residual * coefs[1]`, in the residual’s dtype.
`beam_router_gating_linear`	Router logits `inp @ weight.T (+ bias)` in `router_dtype`; see :class:`_BeamRouterGatingLinear`.
`_fix_group_index_map`	Point `model_param_group_index_map` at the inner optimizer’s parameter order.
`patch_distributed_optimizer_group_index_map`	Monkey patch `DistributedOptimizer.__init__` to apply :func:`_fix_group_index_map` (idempotent).
`_beam_attention_spec`	Rewrite Megatron-Core’s TE `SelfAttention` spec to the Beam gate and SSMax core attention.
`_beam_moe_spec`	Swap `BeamTopKRouter` into Megatron-Core’s MoE spec.
`beam_hybrid_stack_spec`	`HybridStack` spec with Beam’s attention (`*`), dense MLP (`-`) and MoE (`E`) layers.

Data#

BEAM_FEATURE_FIELDS

API#

bridge.models.beam.modeling_beam.beam_logit_softcapping( logits: torch.Tensor, cap: Optional[float], ) → torch.Tensor#

cap * tanh(logits / cap) computed and returned in fp32, exactly as BeamForCausalLM does.

None / 0 disables the cap and returns logits unchanged (in their own dtype).

class bridge.models.beam.modeling_beam.BeamOutputLayer#

Bases: megatron.core.tensor_parallel.ColumnParallelLinear

ColumnParallelLinear whose output is soft-capped in fp32.

Applied to the output layer by BeamModelBuilder.

forward(*args, **kwargs)#: Project, then soft-cap in fp32.

bridge.models.beam.modeling_beam.BEAM_FEATURE_FIELDS#: (‘attention_output_gate_bias’, ‘post_self_attn_layernorm’, ‘post_mlp_layernorm’, ‘depth_scaled_resid…

bridge.models.beam.modeling_beam.validate_beam_config(config) → None#

Reject Beam features on code paths that would silently skip them.

SSMax runs in core attention, sandwich norms and depth-scaled residuals replace the layer’s bias-dropout-add; paths that bypass either would build a different model.

bridge.models.beam.modeling_beam.ssmax_absolute_positions( local_seq_len: int, position_offset: int = 0, cp_rank: int = 0, cp_size: int = 1, device=None, ) → torch.Tensor#

Absolute sequence positions for the tokens this rank actually holds.

Without CP, positions are contiguous. Otherwise, get_batch_on_this_cp_rank splits the sequence into 2 * cp_size chunks and gives rank r chunks (r, 2 * cp_size - r - 1). For CP=2 with 16 tokens::

rank 0 local -> [0 1 2 3][12 13 14 15]
rank 1 local -> [4 5 6 7][ 8  9 10 11]

bridge.models.beam.modeling_beam.ssmax_packed_positions( cu_seqlens: torch.Tensor, position_offset: int = 0, cp_rank: int = 0, cp_size: int = 1, total_tokens: int | None = None, ) → torch.Tensor#

Within-sequence positions for every token of a packed (thd) buffer.

As in RoPE’s thd path, use global cumulative lengths (the padded variant for padded batches). Apply :func:ssmax_absolute_positions’ zigzag layout independently to each sequence under CP.

Fully vectorised; passing total_tokens (the local buffer length) avoids a host sync. Tokens past the last real token (tail padding) are clamped onto the final sequence; they are masked out of attention anyway.

bridge.models.beam.modeling_beam.apply_ssmax_to_query( query: torch.Tensor, position_offset: int = 0, cp_rank: int = 0, cp_size: int = 1, cu_seqlens: torch.Tensor | None = None, ) → torch.Tensor#

Scale query by (1 + (position + 1) / 2048) ** 0.1 for each of its positions.

Parameters:

query – [sq, b, np, hn] normally, [t, np, hn] for packed thd batches.
position_offset – Absolute position of query[0] (non-zero during incremental decoding).
cp_rank – This rank’s index within the context-parallel group.
cp_size – Context-parallel world size; see :func:ssmax_absolute_positions.
cu_seqlens – Global cumulative sequence lengths of a packed batch; when given, positions restart at every sequence (see :func:ssmax_packed_positions).

Returns:

The scaled query, same shape and dtype as the input.

bridge.models.beam.modeling_beam._is_ssmax_layer(config, layer_number: int) → bool#: SSMax applies to the NoPE layers (no_rope_freq is 0-indexed, layer_number 1-indexed).

class bridge.models.beam.modeling_beam._BeamSSMaxCoreAttentionMixin(*args, **kwargs)#

Core attention that applies SSMax to the query of NoPE layers.

Attention calls core attention right after RoPE, on both the plain and the activation-checkpointed path, so scaling here is equivalent to scaling in Attention.forward. Paths that bypass core attention (flash decoding, dynamic batching) are rejected by :func:validate_beam_config and :class:BeamSelfAttention.

Initialization

forward(query, key, value, attention_mask, *args, **kwargs)#: Scale the query on NoPE layers, then run the wrapped core attention.

_beam_apply_ssmax(query, key, packed_seq_params)#

class bridge.models.beam.modeling_beam.BeamTEDotProductAttention(*args, **kwargs)#

Bases: bridge.models.beam.modeling_beam._BeamSSMaxCoreAttentionMixin, megatron.core.extensions.transformer_engine.TEDotProductAttention

Transformer Engine TEDotProductAttention with SSMax on NoPE layers.

Initialization

bridge.models.beam.modeling_beam.apply_output_gate( x: torch.Tensor, gate: torch.Tensor, bias: Optional[torch.Tensor] = None, ) → torch.Tensor#

Gate x with sigmoid(gate [+ bias per head]).

The sigmoid (and the bias add) run in fp32; the result is cast back to x.dtype.

Parameters:

x – attention output, [..., num_heads * head_dim] with heads major.
gate – gate pre-activations with as many elements as x.
bias – optional [num_heads] bias, one scalar per (local) head.

class bridge.models.beam.modeling_beam.AttentionOutputGateBias( config, num_heads_per_partition: int, tp_group: Optional[torch.distributed.ProcessGroup] = None, )#

Bases: megatron.core.transformer.module.MegatronModule

Zero-initialized, TP-sharded per-head bias for the attention output gate.

A dedicated module owns the bias so checkpoint converters can identify it by class name.

Initialization

reset_parameters() → None#: Initialize the bias to zero.

sharded_state_dict( prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int], ...] = (), metadata: Optional[dict] = None, )#: Shard bias along dim 0 across tensor-parallel ranks.

class bridge.models.beam.modeling_beam.BeamSelfAttention(config, *args, **kwargs)#

Bases: megatron.core.transformer.attention.SelfAttention

SelfAttention whose output gate carries a per-head bias.

The bias follows the local query heads, including under KV replication (TP > KV heads).

Initialization

_apply_output_gate(x, gate)#

forward_pre_attn_and_core_attn(*args, **kwargs)#: Reject dynamic-batching inference on SSMax layers: it bypasses core attention.

bridge.models.beam.modeling_beam.depth_scaled_residual_coefs( layer_number: int, ) → Tuple[float, float]#

(sublayer_scale, residual_scale) for the residual add of one hybrid layer; the squares sum to one.

residual_add_idx = layer_number + 1 with HybridStack’s 1-based layer_number, which counts attention and MLP layers separately: the embedding is the first term of the residual stream, so Beam layer i (0-based) adds at 2i + 2 (attention) and 2i + 3 (MLP).

bridge.models.beam.modeling_beam.scaled_residual_add( output: torch.Tensor, residual: torch.Tensor, coefs: Tuple[float, float], ) → torch.Tensor#

output * coefs[0] + residual * coefs[1], in the residual’s dtype.

Scalar coefficients avoid saving the residual for backward, allowing activation offload to release its storage after the add.

class bridge.models.beam.modeling_beam.BeamTransformerLayerSubmodules#

Bases: megatron.core.transformer.transformer_layer.TransformerLayerSubmodules

TransformerLayerSubmodules plus the two sandwich norms.

post_self_attn_layernorm: type#: None

post_mlp_layernorm: type#: None

class bridge.models.beam.modeling_beam.BeamTransformerLayerMixin(

config,

submodules,

layer_number: int = 1,

*args,

**kwargs,

)#

Sandwich norms and depth-scaled residual adds for an attention-only or MLP-only hybrid layer.

The bias-dropout-add of the sublayer the layer owns is replaced by :meth:_beam_bda, which computes bias -> post norm -> dropout -> (depth-scaled) residual add. The absent sublayer keeps its identity pass-through. Everything else of the base layer (offload commits, recompute hooks, viewless outputs) stays in place.

Initialization

_beam_build_post_norm(norm)#

_beam_bda(training: bool, fused: bool, *, is_attention: bool)#: Same call contract as get_bias_dropout_add: returns f(x_with_bias, residual, prob).

class bridge.models.beam.modeling_beam.BeamTransformerLayer(

config,

submodules,

layer_number: int = 1,

*args,

**kwargs,

)#

Bases: bridge.models.beam.modeling_beam.BeamTransformerLayerMixin, megatron.core.transformer.transformer_layer.TransformerLayer

Attention-only hybrid layer (*) with Beam’s sandwich norm and depth-scaled residual.

Initialization

class bridge.models.beam.modeling_beam.BeamMLPLayer(

config,

submodules,

layer_number: int = 1,

*args,

**kwargs,

)#

Bases: bridge.models.beam.modeling_beam.BeamTransformerLayerMixin, megatron.core.ssm.mlp_layer.MLPLayer

Dense MLP hybrid layer (-) with Beam’s sandwich norm and depth-scaled residual.

Initialization

class bridge.models.beam.modeling_beam.BeamMoETransformerLayer(

config,

submodules,

layer_number: int = 1,

*args,

**kwargs,

)#

Bases: bridge.models.beam.modeling_beam.BeamTransformerLayerMixin, megatron.core.transformer.transformer_layer.MoETransformerLayer

MoE hybrid layer (E) with Beam’s sandwich norm and depth-scaled residual.

Initialization

class bridge.models.beam.modeling_beam._BeamRouterGatingLinear#

Bases: torch.autograd.Function

RouterGatingLinearFunction for routers whose weight dtype differs from the input dtype.

Two changes to Megatron-Core’s version: cuBLASLt requires both GEMM operands in router_dtype (fp32 router weights meet bf16 activations; saved activations keep their input dtype), and the TE GEMM reads raw storage and ignores strides, so an expanded (stride-0) grad_output such as logits.sum().backward() produces is made contiguous.

static forward(ctx, inp, weight, bias, router_dtype)#

static backward(ctx, grad_output)#

bridge.models.beam.modeling_beam.beam_router_gating_linear(inp, weight, bias, router_dtype)#: Router logits inp @ weight.T (+ bias) in router_dtype; see :class:_BeamRouterGatingLinear.

class bridge.models.beam.modeling_beam.BeamTopKRouter#

Bases: megatron.core.transformer.moe.router.TopKRouter

TopKRouter that keeps its weight (and bias) in fp32 when moe_router_fp32_params is set.

Parameters marked keep_in_fp32 survive Float16Module and checkpoint loads, so the router reproduces HF’s fp32 expert selection.

reset_parameters()#: Initialize like Router.reset_parameters, casting to fp32 instead of params_dtype.

gating(input: torch.Tensor)#: Router.gating through :func:beam_router_gating_linear.

bridge.models.beam.modeling_beam._fix_group_index_map(optimizer) → None#

Point model_param_group_index_map at the inner optimizer’s parameter order.

DistributedOptimizer builds each inner-optimizer group as fp32 shards followed by float16-derived shards, while model_param_group_index_map follows the original parameter order. With mixed dtypes in one group (fp32 router + bf16 rest) the two orders differ and checkpointed optimizer state lands in the wrong slots.

bridge.models.beam.modeling_beam.patch_distributed_optimizer_group_index_map() → None#: Monkey patch DistributedOptimizer.__init__ to apply :func:_fix_group_index_map (idempotent).

bridge.models.beam.modeling_beam._beam_attention_spec( attention: megatron.core.transformer.spec_utils.ModuleSpec, ) → megatron.core.transformer.spec_utils.ModuleSpec#: Rewrite Megatron-Core’s TE SelfAttention spec to the Beam gate and SSMax core attention.

bridge.models.beam.modeling_beam._beam_moe_spec(mlp)#: Swap BeamTopKRouter into Megatron-Core’s MoE spec.

bridge.models.beam.modeling_beam.beam_hybrid_stack_spec( config, ) → megatron.core.transformer.spec_utils.ModuleSpec#

HybridStack spec with Beam’s attention (*), dense MLP (-) and MoE (E) layers.

The layers reuse the attention, MLP and MoE submodules of Megatron-Core’s Transformer Engine GPT layer spec, split into one sublayer per hybrid layer, with sandwich norms as TE norms. Accepts a TransformerConfig or a model config with a nested transformer.

bridge.models.beam.modeling_beam#

Module Contents#

Classes#

Functions#

Data#

API#

`bridge.models.beam.modeling_beam`#