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

sharded_state_dict(
prefix: str = '',
sharded_offsets: tuple = (),
metadata: dict | None = None,
) → megatron.core.dist_checkpointing.mapping.ShardedStateDict#

Build the sharded state dict, using combined residual keys when loading them.

_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[megatron.core.transformer.residual_connection.ResidualBranchInput[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,
state: megatron.core.transformer.residual_connection.ResidualConnectionState,
recompute_context: megatron.core.transformer.residual_recompute.ResidualStreamRecomputeContext | None = None,
) → torch.Tensor#

Write the mixer update, replaying only nonterminal residual writes.

_pack_two_stage_output(
ssm_output: torch.Tensor,
state: megatron.core.transformer.residual_connection.ResidualConnectionState,
) → tuple[torch.Tensor, torch.Tensor, megatron.core.transformer.residual_connection.ResidualConnectionState]#

Keep WR’s existing two-stage tuple for ShortcutMoE callers.

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.