`bridge.models.beam.beam_builder`#

Model builder for Beam: the Hybrid builder with the Beam stack spec and final-logit soft-capping.

Module Contents#

Build HybridModel stages from the Beam stack spec; soft-cap the output layer on the last one.

class bridge.models.beam.beam_builder.BeamModelBuilder( model_config: megatron.bridge.models.beam.beam_config.BeamModelConfig, )#

Bases: megatron.bridge.models.hybrid.hybrid_builder.HybridModelBuilder

Build HybridModel stages from the Beam stack spec; soft-cap the output layer on the last one.

Initialization

build_model( pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, pre_process: bool | None = None, post_process: bool | None = None, vp_stage: int | None = None, ) → megatron.core.models.hybrid.hybrid_model.HybridModel#: Build one stage; the stage that owns output_layer gets the soft cap.