[0.5/4] Add streamwise wide-residual connection modules (#7599)

Signed-off-by: Kwangjun Ahn <kwangjuna@nvidia.com>

Author: kwangjunahn

megatron/core/optimizer/__init__.py

@@ -101,19 +101,23 @@ def get_standard_config_overrides(config: OptimizerConfig) -> Dict[ParamKey, Par
         Dict[ParamKey, ParamGroupOverride]: standard config overrides.
     """
     config_overrides: Optional[Dict[ParamKey, ParamGroupOverride]] = {}
-    # First, figure out how we are going to do wd skipping. The two main approaches are:
-    #  1. The classic megatron approach of skipping all len 1 and bias parameters.
-    #  2. The Qwen3-Next approach of doing 1, other than qk layernorm parameters.
+    # Select the model-family convention for zero weight decay on vector-like parameters:
+    # the classic rule skips all 1-D parameters and biases, while the Qwen3-Next rule keeps
+    # weight decay on Q/K layernorm parameters. Wide-residual retention controllers intentionally
+    # follow the run's ordinary weight-decay policy rather than this generic vector exemption.
     if config.apply_wd_to_qk_layernorm:
         shape_1_not_qkln_param = ParamWithNamePredicate(
             name="s1_not_qkln",
             fn=lambda param, name: (len(param.shape) == 1 or name.endswith(".bias"))
-            and not ("q_layernorm." in name or "k_layernorm." in name),
+            and not ("q_layernorm." in name or "k_layernorm." in name)
+            and not getattr(param, "is_wide_residual_retention_parameter", False),
         )
         param_wd_mult_key = ParamKey(with_name_predicate=shape_1_not_qkln_param)
     else:
         param_length_1_match = ParamPredicate(
-            name="param_len_1", fn=lambda param: len(param.shape) == 1
+            name="param_len_1_except_wide_residual_retention",
+            fn=lambda param: len(param.shape) == 1
+            and not getattr(param, "is_wide_residual_retention_parameter", False),
         )
         param_wd_mult_key = ParamKey(name="*.bias", predicate=param_length_1_match)
 

megatron/core/transformer/__init__.py

@@ -8,3 +8,4 @@
     TransformerLayer,
     TransformerLayerSubmodules,
 )
+from .wide_residual_config import WideResidualConfig

megatron/core/transformer/transformer_config.py

@@ -32,6 +32,7 @@
     LayerType,
 )
 from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
+from megatron.core.transformer.wide_residual_config import WideResidualConfig
 
 from .._rank_utils import log_single_rank
 from ..fusions.fused_bias_geglu import quick_gelu
@@ -1290,6 +1291,16 @@ class TransformerConfig(ModelParallelConfig):
 
     Must be a positive integer when set."""
 
+    ####################
+    # Wide Residual Configuration
+    ####################
+    wide_residual: Optional[WideResidualConfig] = None
+    """Optional streamwise wide-residual architecture configuration.
+
+    When set, the model carries ``num_streams * hidden_size`` features between
+    layers while attention and MLP branches continue to operate at ``hidden_size``.
+    """
+
     ####################
     # miscellaneous
     ####################

megatron/core/transformer/wide_residual_config.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+from dataclasses import dataclass
+
+
+@dataclass
+class WideResidualConfig:
+    """Configuration for streamwise wide residuals around ordinary-width branches.
+
+    The model carries ``num_streams`` contiguous residual streams, each with
+    ``TransformerConfig.hidden_size`` features. Attention, MLP, and MoE branches
+    continue to operate at the ordinary hidden size.
+    """
+
+    num_streams: int
+    """Number of ordinary-width streams carried by the model; must be greater than one."""
+
+    streamwise_sigmoid_init_scale: float = 0.01
+    """Symmetric initialization spread for streamwise write logits."""
+
+    learned_retention: bool = False
+    """Apply one bounded learned carry factor to each residual stream."""
+
+    retention_init: float = 0.999
+    """Initial retention factor; values near one preserve the initial function."""
+
+    retention_max_forget: float = 0.10
+    """Maximum forget rate in ``1 - max_forget * sigmoid(-logit)``."""
+
+    def __post_init__(self) -> None:
+        if isinstance(self.num_streams, bool) or not isinstance(self.num_streams, int):
+            raise TypeError("wide residual num_streams must be an integer.")
+        if self.num_streams <= 1:
+            raise ValueError(
+                f"wide residual num_streams must be greater than one, got {self.num_streams}."
+            )
+        if self.streamwise_sigmoid_init_scale < 0.0:
+            raise ValueError(
+                "streamwise_sigmoid_init_scale must be non-negative, got "
+                f"{self.streamwise_sigmoid_init_scale}."
+            )
+        if self.learned_retention:
+            if not 0.0 < self.retention_max_forget < 1.0:
+                raise ValueError(
+                    f"retention_max_forget must be in (0, 1), got {self.retention_max_forget}."
+                )
+            if not 1.0 - self.retention_max_forget < self.retention_init < 1.0:
+                raise ValueError(
+                    "retention_init must satisfy 1 - retention_max_forget < "
+                    f"retention_init < 1, got {self.retention_init}."
+                )

megatron/core/transformer/wide_residual_layer.py

@@ -0,0 +1,270 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
+"""Construction-time wide-residual connections and model-boundary helpers."""
+
+from __future__ import annotations
+
+import math
+from typing import Literal
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.residual_connection import (
+    ResidualBranchOutput,
+    ResidualConnection,
+    ResidualConnectionState,
+    ResidualConnectionWriteState,
+)
+from megatron.core.transformer.streamwise_residual_ops import (
+    streamwise_sigmoid_read,
+    streamwise_sigmoid_writeback,
+)
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+_MIN_CONTROLLER_NUMEL = 128
+
+
+def _make_padded_vector(values: Tensor) -> nn.Parameter:
+    """Create a parameter vector padded for small distributed-optimizer shards."""
+
+    if values.ndim != 1:
+        raise ValueError(f"Expected one-dimensional initial values, got {values.shape}.")
+    if values.numel() >= _MIN_CONTROLLER_NUMEL:
+        return nn.Parameter(values.clone())
+    padded = torch.zeros(_MIN_CONTROLLER_NUMEL, dtype=values.dtype)
+    padded[: values.numel()] = values
+    return nn.Parameter(padded)
+
+
+def _active_vector(param: Tensor, length: int) -> Tensor:
+    """Return the active prefix of a potentially padded vector parameter."""
+
+    return param[:length]
+
+
+def _mark_residual_map_parameter(param: nn.Parameter, config: TransformerConfig) -> None:
+    """Attach optimizer and TP-gradient metadata to a replicated residual map."""
+
+    setattr(param, "is_residual_map_parameter", True)
+    setattr(param, "use_muon", False)
+    setattr(param, "sequence_parallel", config.sequence_parallel)
+
+
+def _mark_replicated_tp_parameter(param: nn.Parameter, config: TransformerConfig) -> None:
+    """Mark a replicated controller for the correct TP gradient reduction."""
+
+    setattr(param, "allreduce", True)
+    setattr(param, "tensor_model_parallel", False)
+    setattr(param, "average_gradients_across_tp_domain", not config.sequence_parallel)
+
+
+def _mark_retention_parameter(param: nn.Parameter, config: TransformerConfig) -> None:
+    """Attach optimizer and TP-gradient metadata to a replicated retention parameter."""
+
+    setattr(param, "is_wide_residual_retention_parameter", True)
+    setattr(param, "sequence_parallel", config.sequence_parallel)
+    _mark_replicated_tp_parameter(param, config)
+
+
+class LearnedWideResidualRetention(nn.Module):
+    """Bounded learned retention with one controller per full-width stream."""
+
+    def __init__(
+        self, config: TransformerConfig, layer_number: int, branch_name: str, *, num_streams: int
+    ) -> None:
+        super().__init__()
+        if config.wide_residual is None:
+            raise ValueError("LearnedWideResidualRetention requires wide_residual config.")
+        wr = config.wide_residual
+        if num_streams != wr.num_streams:
+            raise ValueError(
+                "Retention must define one controller per full-width stream, got "
+                f"{num_streams} controllers for {wr.num_streams} streams."
+            )
+        self.layer_number = layer_number
+        self.branch_name = branch_name
+        self.num_streams = num_streams
+        self.max_forget = wr.retention_max_forget
+        forget_ratio = (1.0 - wr.retention_init) / self.max_forget
+        initial_logit = math.log((1.0 - forget_ratio) / forget_ratio)
+        self.retention_logit = _make_padded_vector(
+            torch.full((num_streams,), initial_logit, dtype=torch.float32)
+        )
+        _mark_retention_parameter(self.retention_logit, config)
+
+    def factors(self) -> Tensor:
+        """Return one FP32 retention factor for every full-width stream."""
+
+        logits = _active_vector(self.retention_logit, self.num_streams).float()
+        return 1.0 - self.max_forget * torch.sigmoid(-logits)
+
+    def forward(self, *, return_logits: bool = False) -> Tensor:
+        """Return padded logits or materialized retention factors."""
+
+        return self.retention_logit if return_logits else self.factors()
+
+
+class StreamwiseSigmoidMap(nn.Module):
+    """Positive scalar controller for each contiguous full-width stream."""
+
+    def __init__(self, config: TransformerConfig, *, map_kind: Literal["read", "write"]) -> None:
+        super().__init__()
+        if config.wide_residual is None:
+            raise ValueError("StreamwiseSigmoidMap requires wide_residual config.")
+        wr = config.wide_residual
+
+        self.map_kind = map_kind
+        self.num_streams = wr.num_streams
+
+        if map_kind == "read":
+            initial_factor = 1.0 / self.num_streams
+            initial_logit = math.log(initial_factor / (1.0 - initial_factor))
+            initial_logits = torch.full((self.num_streams,), initial_logit, dtype=torch.float32)
+        elif map_kind == "write":
+            initial_logits = torch.linspace(
+                -wr.streamwise_sigmoid_init_scale,
+                wr.streamwise_sigmoid_init_scale,
+                self.num_streams,
+                dtype=torch.float32,
+            )
+        else:
+            raise ValueError(f"Unsupported streamwise sigmoid map kind: {map_kind!r}.")
+
+        self.logit = _make_padded_vector(initial_logits)
+        _mark_residual_map_parameter(self.logit, config)
+        _mark_replicated_tp_parameter(self.logit, config)
+
+    def factors(self) -> Tensor:
+        """Return one positive factor for every full-width residual stream."""
+
+        factors = torch.sigmoid(_active_vector(self.logit, self.num_streams).float())
+        if self.map_kind == "write":
+            factors = 2.0 * factors
+        return factors
+
+    def forward(self, *, return_logits: bool = False) -> Tensor:
+        """Return padded logits or one positive factor per full-width stream."""
+
+        return self.logit if return_logits else self.factors()
+
+
+class StreamwiseSigmoidWideResidualConnection(ResidualConnection):
+    """Positive streamwise maps around one ordinary-width residual branch."""
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        layer_number: int,
+        branch_name: str,
+        pg_collection: ProcessGroupCollection,
+        name: str | None = None,
+    ) -> None:
+        del pg_collection, name
+        if config.wide_residual is None:
+            raise ValueError(
+                "StreamwiseSigmoidWideResidualConnection requires wide_residual config."
+            )
+        wr = config.wide_residual
+        super().__init__(
+            residual_stream_hidden_size=wr.num_streams * config.hidden_size,
+            branch_hidden_size=config.hidden_size,
+        )
+        self.layer_number = layer_number
+        self.branch_name = branch_name
+        self.num_streams = wr.num_streams
+        self.read_map = StreamwiseSigmoidMap(config, map_kind="read")
+        self.write_map = StreamwiseSigmoidMap(config, map_kind="write")
+        self.retention = (
+            LearnedWideResidualRetention(
+                config, layer_number, branch_name, num_streams=self.num_streams
+            )
+            if wr.learned_retention
+            else None
+        )
+
+    def _read(self, hidden_states: Tensor) -> tuple[Tensor, ResidualConnectionWriteState]:
+        return (
+            streamwise_sigmoid_read(
+                hidden_states, self.read_map(return_logits=True), self.num_streams
+            ),
+            (),
+        )
+
+    def _write(
+        self,
+        branch_output: ResidualBranchOutput,
+        state: ResidualConnectionState,
+        *,
+        dropout_probability: float,
+        training: bool,
+    ) -> Tensor:
+        residual_stream = state[0]
+        if isinstance(branch_output, tuple):
+            branch_update, bias = branch_output
+        else:
+            branch_update, bias = branch_output, None
+
+        branch_update = branch_update.to(dtype=residual_stream.dtype)
+        if bias is not None:
+            branch_update = branch_update + bias.to(
+                device=branch_update.device, dtype=branch_update.dtype
+            )
+        branch_update = F.dropout(branch_update, p=dropout_probability, training=training)
+        retention_logits = (
+            self.retention(return_logits=True) if self.retention is not None else None
+        )
+        return streamwise_sigmoid_writeback(
+            residual_stream,
+            branch_update,
+            self.write_map(return_logits=True),
+            self.num_streams,
+            retention_logits=retention_logits,
+            retention_max_forget=(self.retention.max_forget if self.retention is not None else 0.0),
+        )
+
+
+class StreamwiseSigmoidResidualReadout(nn.Module):
+    """Learned streamwise readout from D' to D before final normalization."""
+
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        if config.wide_residual is None:
+            raise ValueError("StreamwiseSigmoidResidualReadout requires wide_residual config.")
+        wr = config.wide_residual
+        self.input_hidden_size = wr.num_streams * config.hidden_size
+        self.output_hidden_size = config.hidden_size
+        self.exit_map = StreamwiseSigmoidMap(config, map_kind="read")
+
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        """Mix the full-width streams into one backbone-width activation."""
+
+        if hidden_states.shape[-1] != self.input_hidden_size:
+            raise ValueError(
+                "Streamwise sigmoid readout expected hidden size "
+                f"{self.input_hidden_size}, got {hidden_states.shape[-1]}."
+            )
+        return streamwise_sigmoid_read(
+            hidden_states, self.exit_map(return_logits=True), self.exit_map.num_streams
+        )
+
+
+def build_wide_residual_readout(
+    config: TransformerConfig,
+) -> StreamwiseSigmoidResidualReadout | None:
+    """Build the streamwise model-boundary readout when wide residuals are enabled."""
+
+    if config.wide_residual is None:
+        return None
+    return StreamwiseSigmoidResidualReadout(config)
+
+
+def expand_wide_residual_stream(hidden_states: Tensor, num_streams: int) -> Tensor:
+    """Replicate an ordinary-width activation into contiguous residual streams."""
+
+    if num_streams <= 1:
+        raise ValueError(f"num_streams must be greater than one, got {num_streams}.")
+    repeat_shape = [1] * hidden_states.ndim
+    repeat_shape[-1] = num_streams
+    return hidden_states.repeat(*repeat_shape)