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 | 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.

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.