bridge.models.mamba.mamba_provider

Module Contents

Classes

MambaModelProvider

Backward-compatible wrapper around :class:HybridModelProvider.

Functions

modelopt_mamba_stack_spec	Backward-compatible alias for modelopt_hybrid_stack_spec.
transformer_engine_mamba_stack_spec	Backward-compatible alias for transformer_engine_hybrid_stack_spec.
get_default_mamba_stack_spec	Backward-compatible alias for get_default_hybrid_stack_spec.

API

class bridge.models.mamba.mamba_provider.MambaModelProvider

Bases: megatron.bridge.models.hybrid.hybrid_provider.HybridModelProvider

Backward-compatible wrapper around :class:HybridModelProvider.

mamba_stack_spec: megatron.core.transformer.ModuleSpec | Callable[[], megatron.core.transformer.ModuleSpec] | Callable[[bridge.models.mamba.mamba_provider.MambaModelProvider], megatron.core.transformer.ModuleSpec] | None

None

__post_init__() → None

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

to_cfg_dict() → dict[str, Any]

Serialize without the deprecated mamba_stack_spec field.

bridge.models.mamba.mamba_provider.modelopt_mamba_stack_spec(

config: MambaModelProvider | None = None,

) → megatron.core.transformer.ModuleSpec

Backward-compatible alias for modelopt_hybrid_stack_spec.

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

Backward-compatible alias for transformer_engine_hybrid_stack_spec.

bridge.models.mamba.mamba_provider.get_default_mamba_stack_spec(

config: bridge.models.mamba.mamba_provider.MambaModelProvider,

) → megatron.core.transformer.ModuleSpec

Backward-compatible alias for get_default_hybrid_stack_spec.