[Refactor] Select wide-residual operation at construction (#7952)

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>

Author: wujingyue

megatron/core/models/hybrid/shortcut_block.py

@@ -193,14 +193,14 @@ def __init__(
                 layer_number=moe_layer.layer_number,
                 branch_name="shortcut_routed",
             )
-            outer_connection = moe_layer._get_mlp_residual_connection()
-            if outer_connection is None:
+            outer_read = moe_layer._get_mlp_residual_connection(operation="read")
+            if outer_read is None:
                 raise ValueError(
                     "A wide-residual shortcut read requires an outer MoE residual connection."
                 )
             if (
                 self.shortcut_residual_read.residual_stream_hidden_size
-                != outer_connection.residual_stream_hidden_size
+                != outer_read.residual_stream_hidden_size
             ):
                 raise ValueError(
                     "The shortcut read and outer MoE residual connection must consume the "

megatron/core/ssm/wide_residual_mamba_layer.py

@@ -13,7 +13,10 @@
     checkpoint_residual_write,
 )
 from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.transformer.wide_residual_layer import StreamwiseSigmoidWideResidualConnection
+from megatron.core.transformer.wide_residual_layer import (
+    StreamwiseSigmoidWideResidualConnection,
+    load_residual_state,
+)
 from megatron.core.typed_torch import apply_module
 
 
@@ -55,33 +58,33 @@ def __init__(
                 "returns its own residual."
             )
 
-        self.residual_connection = StreamwiseSigmoidWideResidualConnection(
-            config=self.config,
-            layer_number=self.layer_number,
-            branch_name="mamba",
-            pg_collection=pg_collection,
-            name=(name + ".residual_connection") if name is not None else None,
+        self.residual_read = StreamwiseSigmoidWideResidualConnection(
+            config, self.layer_number, "mamba", operation="read"
+        )
+        self.residual_write = StreamwiseSigmoidWideResidualConnection(
+            config, self.layer_number, "mamba", operation="write"
         )
+        self.register_load_state_dict_pre_hook(load_residual_state)
         self.residual_stream_hidden_size = (
             config.wide_residual.num_streams * self.config.hidden_size
         )
-        if self.residual_connection.residual_stream_hidden_size != self.residual_stream_hidden_size:
+        if self.residual_read.residual_stream_hidden_size != self.residual_stream_hidden_size:
             raise ValueError(
                 "The wide-residual Mamba connection must carry num_streams * hidden_size "
                 f"features, expected {self.residual_stream_hidden_size}, got "
-                f"{self.residual_connection.residual_stream_hidden_size}."
+                f"{self.residual_read.residual_stream_hidden_size}."
             )
-        if self.residual_connection.branch_hidden_size != self.config.hidden_size:
+        if self.residual_read.branch_hidden_size != self.config.hidden_size:
             raise ValueError(
                 "The wide-residual Mamba connection must produce "
                 f"hidden_size={self.config.hidden_size}, got "
-                f"{self.residual_connection.branch_hidden_size}."
+                f"{self.residual_read.branch_hidden_size}."
             )
 
     def _get_residual_connection(self):
         """Return the connection surrounding the Mamba mixer."""
 
-        return self.residual_connection
+        return self.residual_read, self.residual_write
 
     def _prepare_mixer_state(
         self,
@@ -92,21 +95,20 @@ def _prepare_mixer_state(
 
         recompute_context = residual_stream_recompute_context
         if recompute_context is None:
-            hidden_states, connection_state = apply_module(self.residual_connection)(
+            hidden_states, connection_state = apply_module(self.residual_read)(
                 hidden_states,
-                operation="read",
                 fp32_residual_connection=self.config.fp32_residual_connection,
                 branch_input_dtype=self.config.params_dtype,
             )
         else:
             hidden_states, connection_state = checkpoint_residual_read(
-                self.residual_connection,
+                self.residual_read,
                 hidden_states,
                 recompute_context,
                 fp32_residual_connection=self.config.fp32_residual_connection,
                 branch_input_dtype=self.config.params_dtype,
             )
-        residual = self.residual_connection.residual_stream(connection_state)
+        residual = self.residual_read.residual_stream(connection_state)
 
         hidden_states = hidden_states.to(dtype=self.config.params_dtype)
         if recompute_context is not None and not isinstance(self.norm, IdentityOp):
@@ -128,17 +130,16 @@ def _apply_mixer_update(
             raise RuntimeError("Missing state for the Mamba residual connection.")
         if recompute_context is not None and not recompute_context.is_block_end:
             return checkpoint_residual_write(
-                self.residual_connection,
+                self.residual_write,
                 mixer_out_with_bias,
                 connection_state,
                 recompute_context,
                 dropout_probability=self.hidden_dropout,
                 training=self.training,
             )
         with self.bias_dropout_add_exec_handler():
-            return apply_module(self.residual_connection)(
+            return apply_module(self.residual_write)(
                 mixer_out_with_bias,
-                operation="write",
                 state=connection_state,
                 dropout_probability=self.hidden_dropout,
                 training=self.training,

megatron/core/transformer/residual_connection.py

@@ -21,10 +21,21 @@ class ResidualConnection(nn.Module, ABC):
     ``forward`` validates the common contract and retains the incoming residual
     stream as the first tensor in ``ResidualConnectionState``. Concrete connections
     own bias, dropout, mapping, and residual-update semantics.
+
+    ``operation`` fixes an instance to one operation at construction.
     """
 
-    def __init__(self, residual_stream_hidden_size: int, branch_hidden_size: int):
+    def __init__(
+        self,
+        residual_stream_hidden_size: int,
+        branch_hidden_size: int,
+        *,
+        operation: ResidualConnectionOperation,
+    ):
         super().__init__()
+        if operation not in ("read", "write"):
+            raise ValueError(f"Unsupported residual connection operation: {operation}.")
+        self.operation = operation
         if residual_stream_hidden_size <= 0:
             raise ValueError("Residual-stream hidden size must be positive.")
         if branch_hidden_size <= 0:
@@ -49,7 +60,6 @@ def forward(
         self,
         value: Tensor,
         *,
-        operation: Literal["read"],
         fp32_residual_connection: bool = False,
         branch_input_dtype: torch.dtype | None = None,
     ) -> tuple[Tensor, ResidualConnectionState]: ...
@@ -59,7 +69,6 @@ def forward(
         self,
         value: ResidualBranchOutput,
         *,
-        operation: Literal["write"],
         state: ResidualConnectionState,
         dropout_probability: float,
         training: bool,
@@ -69,7 +78,6 @@ def forward(
         self,
         value: ResidualBranchOutput,
         *,
-        operation: ResidualConnectionOperation,
         state: ResidualConnectionState | None = None,
         fp32_residual_connection: bool = False,
         branch_input_dtype: torch.dtype | None = None,
@@ -82,37 +90,39 @@ def forward(
         dtype; otherwise compatible connections may produce that dtype directly in a fused read.
         """
 
-        if operation == "read":
-            if not torch.is_tensor(value):
-                raise TypeError("Residual connection read expects a tensor.")
-            if state is not None or dropout_probability is not None or training is not None:
-                raise TypeError("Residual connection read received write-only arguments.")
-            if branch_input_dtype is not None and (
-                not isinstance(branch_input_dtype, torch.dtype)
-                or not branch_input_dtype.is_floating_point
-            ):
-                raise TypeError(
-                    "branch_input_dtype must be a floating-point torch.dtype or None, got "
-                    f"{branch_input_dtype!r}."
+        match self.operation:
+            case "read":
+                if not torch.is_tensor(value):
+                    raise TypeError("Residual connection read expects a tensor.")
+                if state is not None or dropout_probability is not None or training is not None:
+                    raise TypeError("Residual connection read received write-only arguments.")
+                if branch_input_dtype is not None and (
+                    not isinstance(branch_input_dtype, torch.dtype)
+                    or not branch_input_dtype.is_floating_point
+                ):
+                    raise TypeError(
+                        "branch_input_dtype must be a floating-point torch.dtype or None, got "
+                        f"{branch_input_dtype!r}."
+                    )
+                return self._read_with_validation(
+                    value,
+                    fp32_residual_connection=fp32_residual_connection,
+                    branch_input_dtype=branch_input_dtype,
                 )
-            return self._read_with_validation(
-                value,
-                fp32_residual_connection=fp32_residual_connection,
-                branch_input_dtype=branch_input_dtype,
-            )
-        if operation == "write":
-            if fp32_residual_connection or branch_input_dtype is not None:
-                raise TypeError("Residual connection write received read-only arguments.")
-            if state is None:
-                raise TypeError("Residual connection write requires connection state.")
-            if dropout_probability is None or training is None:
-                raise TypeError(
-                    "Residual connection write requires dropout_probability and training."
+            case "write":
+                if fp32_residual_connection or branch_input_dtype is not None:
+                    raise TypeError("Residual connection write received read-only arguments.")
+                if state is None:
+                    raise TypeError("Residual connection write requires connection state.")
+                if dropout_probability is None or training is None:
+                    raise TypeError(
+                        "Residual connection write requires dropout_probability and training."
+                    )
+                return self._write_with_validation(
+                    value, state, dropout_probability=dropout_probability, training=training
                 )
-            return self._write_with_validation(
-                value, state, dropout_probability=dropout_probability, training=training
-            )
-        raise ValueError(f"Unsupported residual connection operation: {operation}.")
+            case _:
+                raise ValueError(f"Unsupported residual connection operation: {self.operation}.")
 
     def _read_with_validation(
         self,

megatron/core/transformer/residual_recompute.py

@@ -119,10 +119,7 @@ def checkpoint_residual_read(
 
     def run_read(stream: Tensor) -> tuple[Tensor, ...]:
         branch_input, state = apply_module(connection)(
-            stream,
-            operation="read",
-            fp32_residual_connection=False,
-            branch_input_dtype=branch_input_dtype,
+            stream, fp32_residual_connection=False, branch_input_dtype=branch_input_dtype
         )
         if state[0].shape != stream.shape:
             raise ValueError("Residual connection read returned an incompatible carried stream.")
@@ -167,7 +164,6 @@ def run_write(
         value = (branch_update, branch_bias) if return_tuple else branch_update
         return apply_module(connection)(
             value,
-            operation="write",
             state=connection_state,
             dropout_probability=dropout_probability,
             training=training,

megatron/core/transformer/transformer_layer.py

@@ -33,6 +33,10 @@
 from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
 from megatron.core.transformer.mlp import MLP
 from megatron.core.transformer.module import GraphableMegatronModule, TwoStageAttentionLayer
+from megatron.core.transformer.residual_connection import (
+    ResidualConnection,
+    ResidualConnectionOperation,
+)
 from megatron.core.transformer.residual_recompute import (
     ResidualStreamRecomputeContext,
     checkpoint_residual_read,
@@ -918,13 +922,17 @@ def forward_post_core_attn(
             residual_stream_recompute_context=residual_stream_recompute_context,
         )
 
-    def _get_self_attention_residual_connection(self):
-        """Return an optional architecture-owned self-attention connection."""
+    def _get_self_attention_residual_connection(
+        self, operation: ResidualConnectionOperation
+    ) -> ResidualConnection | None:
+        """Return the optional architecture-owned self-attention reader or writer."""
 
         return None
 
-    def _get_mlp_residual_connection(self):
-        """Return an optional architecture-owned MLP connection."""
+    def _get_mlp_residual_connection(
+        self, operation: ResidualConnectionOperation
+    ) -> ResidualConnection | None:
+        """Return the optional architecture-owned MLP reader or writer."""
 
         return None
 
@@ -948,7 +956,7 @@ def _run_input_layernorm(
             ``_apply_self_attn_bda_step``. A configured residual connection uses
             it for the corresponding write state; otherwise base returns ``()``.
         """
-        residual_connection = self._get_self_attention_residual_connection()
+        residual_connection = self._get_self_attention_residual_connection(operation="read")
         recompute_context = (
             residual_stream_recompute_context if residual_connection is not None else None
         )
@@ -957,7 +965,6 @@ def _run_input_layernorm(
             if recompute_context is None:
                 hidden_states, connection_state = apply_module(residual_connection)(
                     hidden_states,
-                    operation="read",
                     fp32_residual_connection=self.config.fp32_residual_connection,
                     branch_input_dtype=self.config.params_dtype,
                 )
@@ -1033,7 +1040,7 @@ def _apply_self_attn_bda_step(
         using_fused_tp_inference_kernel = (
             InferenceMode.is_active() and self.config.inference_fuse_tp_communication
         )
-        residual_connection = self._get_self_attention_residual_connection()
+        residual_connection = self._get_self_attention_residual_connection(operation="write")
         recompute_context = (
             residual_stream_recompute_context if residual_connection is not None else None
         )
@@ -1051,7 +1058,7 @@ def _apply_self_attn_bda_step(
             is_terminal_write = bool(
                 recompute_context is not None
                 and recompute_context.is_block_end
-                and self._get_mlp_residual_connection() is None
+                and self._get_mlp_residual_connection(operation="write") is None
             )
             if recompute_context is not None and not is_terminal_write:
                 hidden_states = checkpoint_residual_write(
@@ -1066,7 +1073,6 @@ def _apply_self_attn_bda_step(
                 with self.bias_dropout_add_exec_handler():
                     hidden_states = apply_module(residual_connection)(
                         attention_output_with_bias,
-                        operation="write",
                         state=attn_state,
                         dropout_probability=self.hidden_dropout,
                         training=self.training,
@@ -1254,7 +1260,7 @@ def _pre_mlp_layernorm_and_residual(
             through to ``_apply_mlp_bda_step``. A configured residual connection
             uses it for the corresponding write state; otherwise base returns ``()``.
         """
-        residual_connection = self._get_mlp_residual_connection()
+        residual_connection = self._get_mlp_residual_connection(operation="read")
         recompute_context = (
             residual_stream_recompute_context if residual_connection is not None else None
         )
@@ -1263,7 +1269,6 @@ def _pre_mlp_layernorm_and_residual(
             if recompute_context is None:
                 hidden_states, connection_state = apply_module(residual_connection)(
                     hidden_states,
-                    operation="read",
                     fp32_residual_connection=self.config.fp32_residual_connection,
                     branch_input_dtype=self.config.params_dtype,
                 )
@@ -1578,7 +1583,7 @@ def _apply_mlp_bda_step(
             InferenceMode.is_active() and self.config.inference_fuse_tp_communication
         )
 
-        residual_connection = self._get_mlp_residual_connection()
+        residual_connection = self._get_mlp_residual_connection(operation="write")
         recompute_context = (
             residual_stream_recompute_context if residual_connection is not None else None
         )
@@ -1613,7 +1618,6 @@ def _apply_mlp_bda_step(
                 with self.bias_dropout_add_exec_handler():
                     hidden_states = apply_module(residual_connection)(
                         mlp_output_with_bias,
-                        operation="write",
                         state=mlp_state,
                         dropout_probability=self.hidden_dropout,
                         training=self.training,