`bridge.models.beam.beam_bridge`#

Beam config and checkpoint conversion.

Beam uses grouped sigmoid routing, sandwich norms, depth-scaled residuals, full-width attention gates, SSMax on NoPE layers, and fp32 logit soft-capping. Router parameters remain fp32; export with weight_dtype=None preserves them. The expert-bias update rate defaults to zero for checkpoint fine-tuning. The modules Megatron-Core lacks for these features live in :mod:.modeling_beam.

The Megatron model is a HybridModel with two hybrid layers per Beam layer i: attention at decoder.layers.{2i} and the dense MLP or MoE at decoder.layers.{2i+1}.

Module Contents#

Classes#

BeamBridge

Bridge between HuggingFace BeamForCausalLM and Megatron-Core HybridModel.

Functions#

`_per_layer_types`	Read a per-layer type list (`layer_types` / `mlp_layer_types`) and validate it.
`_moe_layer_pattern`	Return 1 for sparse layers and 0 for dense layers.
`_rope_settings`	`(no_rope_freq, rotary_base)` from `rope_parameters` keyed by layer type.
`_on_attention_layers`	Spread per-Beam-layer flags over the hybrid layers: set on attention (even) layers, 0 on MLP layers.
`_beam_settings`	Megatron config settings for a `BeamConfig`.
`_as_layer_list`	Expand a Megatron `int \| list` per-layer knob into a list of length `num_layers`.
`_attention_layer_mappings`	Parameters of one attention hybrid layer.
`_dense_mlp_layer_mappings`	Parameters of one dense MLP hybrid layer (the pre-MLP norm is fused into `linear_fc1`).
`_moe_layer_mappings`	Parameters of one MoE hybrid layer.

Data#

`LAYER_TYPES`
`MLP_LAYER_TYPES`
`_DENSE_EXPORT_DROPS`

API#

bridge.models.beam.beam_bridge.LAYER_TYPES#: (‘sliding_attention’, ‘full_attention’)

bridge.models.beam.beam_bridge.MLP_LAYER_TYPES#: (‘dense’, ‘sparse’)

bridge.models.beam.beam_bridge._per_layer_types( hf_config, name: str, allowed: Collection[str], ) → list[str]#: Read a per-layer type list (layer_types / mlp_layer_types) and validate it.

bridge.models.beam.beam_bridge._moe_layer_pattern(hf_config) → list[int]#: Return 1 for sparse layers and 0 for dense layers.

bridge.models.beam.beam_bridge._rope_settings(hf_config) → tuple[list[int] | None, float | None]#

(no_rope_freq, rotary_base) from rope_parameters keyed by layer type.

HF BeamModel computes RoPE for a layer type whose rope_parameters entry is a dict and runs NoPE + SSMax for a layer type whose entry is None. Megatron marks NoPE per layer (no_rope_freq[i] truthy = skip RoPE) and has a single rotary_base, so every RoPE layer type must use plain RoPE with the same theta.

no_rope_freq is None when every layer uses RoPE (core’s default); rotary_base is None when no layer does.

bridge.models.beam.beam_bridge._on_attention_layers(flags: list[int]) → list[int]#: Spread per-Beam-layer flags over the hybrid layers: set on attention (even) layers, 0 on MLP layers.

bridge.models.beam.beam_bridge._beam_settings(hf_config) → dict[str, Any]#: Megatron config settings for a BeamConfig.

bridge.models.beam.beam_bridge._DENSE_EXPORT_DROPS#: (‘num_experts’, ‘moe_intermediate_size’, ‘num_experts_per_tok’, ‘norm_topk_prob’, ‘n_group’, ‘topk_g…

bridge.models.beam.beam_bridge._as_layer_list(value, num_layers: int, *, int_rule) → list[int]#: Expand a Megatron int | list per-layer knob into a list of length num_layers.

class bridge.models.beam.beam_bridge.BeamBridge#

Bases: megatron.bridge.models.conversion.model_bridge.MegatronModelBridge

Bridge between HuggingFace BeamForCausalLM and Megatron-Core HybridModel.

.. rubric:: Example

from megatron.bridge import AutoBridge bridge = AutoBridge.from_hf_pretrained(””) model_config = bridge.get_model_config() model = bridge.get_model(model_config)

The builder path is preferred; to_megatron_provider() remains available for legacy consumers.

MODEL_CONFIG_CLASS#: None

USE_MODEL_CONFIG_FOR_CONVERSION#: True

_adapter_target_modules( adapter_weight_names: Iterable[str], ) → list[str]#

Use exact paths: Beam’s attention and MLP both contain gate_proj.

This also matters for MLP-only adapters: a bare gate_proj would silently attach an unexported adapter to every attention gate and other layers.

_is_fused_qkv(hf_weight_names: Iterable[str]) → bool#: Recognize Beam’s four attention projections, including its independent gate.

_infer_qkv_projection_from_name(hf_name: str) → str | None#: Keep attention gate_proj distinct from the MLP gate.

_split_qkv_linear_out_weight( megatron_model: megatron.core.models.hybrid.hybrid_model.HybridModel | list[megatron.core.models.hybrid.hybrid_model.HybridModel], linear_out_weight: torch.Tensor, ) → dict[str, torch.Tensor]#

Split fused LoRA B along the same Q/gate/K/V rows as the base weight.

The trailing dimension is the adapter rank, not the model hidden size. The export and merged-weight paths both use this hook; LoRA A is shared.

FP8_WEIGHT_SCALE_SUFFIX#: ‘_scale’

static get_hf_import_param_names( hf_param: str | dict[str, str], available_hf_param_names: set[str] | None = None, ) → tuple[str, ...]#: Also declare the weight_scale companions so incremental imports wait for them.

maybe_modify_loaded_hf_weight( hf_param: str | dict[str, str], hf_state_dict: Mapping[str, torch.Tensor], ) → torch.Tensor | dict[str, torch.Tensor]#: Dequantize per-output-channel FP8 weights on the fly; bf16 tensors pass through.

classmethod _maybe_dequantize_fp8( weight: torch.Tensor, param_name: str, hf_state_dict: Mapping[str, torch.Tensor], ) → torch.Tensor#: fp8_weight * weight_scale in fp32 (one rounding when the mapping casts to the param dtype).

hf_config_to_model_config_kwargs(hf_config) → dict[str, Any]#: Flat kwargs for :class:BeamModelConfig (builder-backed path).

classmethod megatron_to_hf_config(provider) → dict#

Invert the config mapping so an exported checkpoint loads as BeamForCausalLM.

provider is a :class:BeamModelConfig, which proxies attribute access to its nested transformer config. Its per-layer settings count hybrid layers; Beam layer i reads them from hybrid layers 2i (attention) and 2i + 1 (MLP).

provider_bridge(hf_pretrained)#: Build a compatibility provider from the same config mapping as the builder.

mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#

Map every Beam layer onto its two hybrid layers.

Beam layer i (model.layers.{i}) is attention at decoder.layers.{2i} and a dense MLP or MoE, per mlp_layer_types, at decoder.layers.{2i+1}.

bridge.models.beam.beam_bridge._attention_layer_mappings(layer: str, hf: str) → list#: Parameters of one attention hybrid layer.

bridge.models.beam.beam_bridge._dense_mlp_layer_mappings(layer: str, hf: str) → list#: Parameters of one dense MLP hybrid layer (the pre-MLP norm is fused into linear_fc1).

bridge.models.beam.beam_bridge._moe_layer_mappings(layer: str, hf: str) → list#: Parameters of one MoE hybrid layer.

bridge.models.beam.beam_bridge#