bridge.models.beam.beam_mappings

Beam full-width attention-gate packing.

Each query group contains [Q, gate, K, V]. Q and gate each have heads_per_group * head_dim rows; K and V each have head_dim rows. TP/PP communication uses the inherited AutoMapping delegate.

Module Contents

Classes

BeamQKVGateMapping

Fuse Q/K/V and a full-width attention gate for Beam.

Functions

_geometry	Return (head_num, num_query_groups, heads_per_group, head_size).
merge_qkv_gate_weights	Interleave Q, K, V and a full-width gate into Megatron’s linear_qkv.
split_qkv_gate_weights	Inverse of :func:merge_qkv_gate_weights; returns (q, k, v, g).

API

bridge.models.beam.beam_mappings._geometry(config) → Tuple[int, int, int, int]

Return (head_num, num_query_groups, heads_per_group, head_size).

bridge.models.beam.beam_mappings.merge_qkv_gate_weights(config, q, k, v, g) → torch.Tensor

Interleave Q, K, V and a full-width gate into Megatron’s linear_qkv.

bridge.models.beam.beam_mappings.split_qkv_gate_weights(config, qkv: torch.Tensor)

Inverse of :func:merge_qkv_gate_weights; returns (q, k, v, g).

class bridge.models.beam.beam_mappings.BeamQKVGateMapping

Bases: megatron.bridge.models.conversion.param_mapping.QKVGMapping

Fuse Q/K/V and a full-width attention gate for Beam.

hf_to_megatron(

hf_weights: dict[str, torch.Tensor],

megatron_module: torch.nn.Module,

) → torch.Tensor

Pack the full-width gate with Q/K/V and distribute the fused weight.

megatron_to_hf(

megatron_weights: torch.Tensor | None,

megatron_module: torch.nn.Module | None,

) → dict[str, torch.Tensor]

Gather the fused weight and split Beam’s four attention projections.