`bridge.models.hybrid.hybrid_builder`#

Module Contents#

Classes#

`HybridModelConfig`	Bridge Hybrid config with transitional output-logit dtype support.
`HybridModelBuilder`	Bridge Hybrid builder that prevents silent fallback on older MCore.

Functions#

`transformer_engine_hybrid_stack_spec`	Return the default Hybrid stack spec with Transformer Engine layers.
`modelopt_hybrid_stack_spec`	Hybrid stack specification for quantization with ModelOpt.
`get_default_hybrid_stack_spec`	Determine the most appropriate Hybrid stack specification based on configuration.

Data#

__all__

API#

class bridge.models.hybrid.hybrid_builder.HybridModelConfig#

Bases: megatron.training.models.hybrid.HybridModelConfig

Bridge Hybrid config with transitional output-logit dtype support.

builder: ClassVar[str]#: ‘megatron.bridge.models.hybrid.HybridModelBuilder’

logit_dtype: torch.dtype | None#: None

class bridge.models.hybrid.hybrid_builder.HybridModelBuilder#

Bases: megatron.training.models.hybrid.HybridModelBuilder

Bridge Hybrid builder that prevents silent fallback on older MCore.

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 a Hybrid stage after validating MCore logit-dtype support.

bridge.models.hybrid.hybrid_builder.transformer_engine_hybrid_stack_spec() → megatron.core.transformer.ModuleSpec#

Return the default Hybrid stack spec with Transformer Engine layers.

This is a named function (not a lambda) to allow proper serialization and reconstruction from checkpoints. Named functions can be imported via their module path, unlike lambdas.

Returns:: Default Hybrid stack specification from megatron.core.

bridge.models.hybrid.hybrid_builder.modelopt_hybrid_stack_spec( config: HybridModelConfig | None = None, ) → megatron.core.transformer.ModuleSpec#

Hybrid stack specification for quantization with ModelOpt.

Uses Norm instead of TENorm and ColumnParallelLinear/RowParallelLinear instead of TE layers to enable proper quantizer insertion by ModelOpt.

Parameters:: config – Optional Hybrid configuration object.
Returns:: Module specification for quantization-ready Hybrid stack.

bridge.models.hybrid.hybrid_builder.get_default_hybrid_stack_spec( config: bridge.models.hybrid.hybrid_builder.HybridModelConfig, ) → megatron.core.transformer.ModuleSpec#

Determine the most appropriate Hybrid stack specification based on configuration.

Parameters:: config – Hybrid configuration object.
Returns:: Appropriate module specification based on config.

bridge.models.hybrid.hybrid_builder.__all__#: [‘HybridModelBuilder’, ‘HybridModelConfig’, ‘get_default_hybrid_stack_spec’, ‘modelopt_hybrid_stack_…

bridge.models.hybrid.hybrid_builder#