bridge.models.mamba.mamba_builder#

Module Contents#

Classes#

MambaModelConfig

Backward-compatible wrapper around :class:HybridModelConfig.

MambaModelBuilder

Backward-compatible wrapper around :class:HybridModelBuilder.

Functions#

transformer_engine_mamba_stack_spec

Backward-compatible alias for transformer_engine_hybrid_stack_spec.

modelopt_mamba_stack_spec

Backward-compatible alias for modelopt_hybrid_stack_spec.

get_default_mamba_stack_spec

Backward-compatible alias for get_default_hybrid_stack_spec.

API#

class bridge.models.mamba.mamba_builder.MambaModelConfig#

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

Backward-compatible wrapper around :class:HybridModelConfig.

builder: ClassVar[str]#

‘megatron.bridge.models.mamba.MambaModelBuilder’

mamba_stack_spec: megatron.core.transformer.ModuleSpec | Callable[[], megatron.core.transformer.ModuleSpec] | Callable[[bridge.models.mamba.mamba_builder.MambaModelConfig], megatron.core.transformer.ModuleSpec] | None#

None

__post_init__() → None#

Normalize the deprecated Mamba stack-spec field to the Hybrid field.

as_dict() → dict[str, Any]#

Serialize without the deprecated mamba_stack_spec field.

_to_dict() → dict[str, Any]#

Bridge config serialization hook.

class bridge.models.mamba.mamba_builder.MambaModelBuilder#

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

Backward-compatible wrapper around :class:HybridModelBuilder.

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 model while resolving legacy Mamba stack-spec factories.

Parameters:

  • pg_collection – Process groups used to construct the model.

  • pre_process – Whether to include input processing on this stage.

  • post_process – Whether to include output processing on this stage.

  • vp_stage – Virtual pipeline stage to construct.

Returns:

The constructed Hybrid model.

bridge.models.mamba.mamba_builder.transformer_engine_mamba_stack_spec() → megatron.core.transformer.ModuleSpec#

Backward-compatible alias for transformer_engine_hybrid_stack_spec.

bridge.models.mamba.mamba_builder.modelopt_mamba_stack_spec(
config: MambaModelConfig | None = None,
) → megatron.core.transformer.ModuleSpec#

Backward-compatible alias for modelopt_hybrid_stack_spec.

bridge.models.mamba.mamba_builder.get_default_mamba_stack_spec(
config: bridge.models.mamba.mamba_builder.MambaModelConfig,
) → megatron.core.transformer.ModuleSpec#

Backward-compatible alias for get_default_hybrid_stack_spec.