bridge.models.hybrid.hybrid_builder

Module Contents

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

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: megatron.training.models.hybrid.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_…