core.ssm.wide_residual_mamba_layer#
Mamba layer specialization for streamwise wide-residual connections.
Module Contents#
Classes#
Mamba layer carrying a wide stream around its ordinary-width mixer. |
API#
- class core.ssm.wide_residual_mamba_layer.WideResidualMambaLayer(
- config: megatron.core.transformer.transformer_config.TransformerConfig,
- submodules: megatron.core.ssm.mamba_layer.MambaLayerSubmodules,
- layer_number: int = 1,
- pg_collection: megatron.core.process_groups_config.ProcessGroupCollection = None,
- pp_layer_offset: int = 0,
- name: str | None = None,
Bases:
megatron.core.ssm.mamba_layer.MambaLayerMamba layer carrying a wide stream around its ordinary-width mixer.
Initialization
Initialize Mamba Layer.
- Parameters:
name (str | None) – module instance name passed top-down from its paranet module
- supports_wide_residual_connections: bool#
True
- _get_residual_connection()#
Return the connection surrounding the Mamba mixer.