Is this page helpful?

core.ssm.wide_residual_mamba_layer#

Mamba layer specialization for streamwise wide-residual connections.

Module Contents#

Classes#

WideResidualMambaLayer

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,
is_mtp_layer: bool = False,
)#

Bases: megatron.core.ssm.mamba_layer.MambaLayer

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

_prepare_mixer_state(
hidden_states: torch.Tensor,
residual_stream_recompute_context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext | None = None,
) → tuple[torch.Tensor, torch.Tensor, megatron.core.transformer.residual_connection.ResidualConnectionState, megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext | None]#

Read the wide stream and optionally replay its connected input normalization.

_apply_mixer_update(
mixer_out_with_bias,
residual: torch.Tensor,
connection_state: megatron.core.transformer.residual_connection.ResidualConnectionState | None = None,
recompute_context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext | None = None,
) → torch.Tensor#

Write the mixer update, replaying only nonterminal residual writes.

forward_post_core_attn(
ssm_output: torch.Tensor,
residual: torch.Tensor,
connection_state: megatron.core.transformer.residual_connection.ResidualConnectionState,
*,
residual_stream_recompute_context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext | None = None,
) → torch.Tensor#

Project the mixer output and write it to the saved wide-residual stream.