| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """HyperConnection (Gated Residual) utilities. |
| 4 | |
| 5 | Implements the HyperConnection residual scheme proposed in |
| 6 | "HyperConnections" (https://arxiv.org/abs/2409.19606). |
| 7 | |
| 8 | The two concrete variants are: |
| 9 | - ``HyperConnectionBase`` - simple average pooling across hc_count parallel |
| 10 | streams (equivalent to hyperconnection_average). |
| 11 | - ``GatedResidual`` - learnable low-rank gated mixing and injection |
| 12 | (gated_residual). |
| 13 | |
| 14 | Hidden states between layers have shape ``[..., HC*HS]`` with HS inner |
| 15 | (HC outer, HS inner — checkpoint-native layout). The local torch |
| 16 | implementation consumes the hyper input viewed as ``[..., HC, HS]``. |
| 17 | |
| 18 | Typical usage inside a transformer decoder layer:: |
| 19 | |
| 20 | self.attn_hc = GatedResidual(hc_config, role="attn") |
| 21 | self.mlp_hc = GatedResidual(hc_config, role="mlp") |
| 22 | |
| 23 | hidden_states, residual = self.attn_hc.mix(hidden_states) |
| 24 | hidden_states = attention(hidden_states) |
| 25 | hidden_states = self.attn_hc.combine(hidden_states, residual) |
| 26 | |
| 27 | hidden_states, residual = self.mlp_hc.mix(hidden_states) |
| 28 | hidden_states = mlp(hidden_states) |
| 29 | hidden_states = self.mlp_hc.combine(hidden_states, residual) |
| 30 | """ |
| 31 | |
| 32 | from dataclasses import dataclass |
| 33 | |
| 34 | import torch |
| 35 | import torch.nn.functional as F |
| 36 | from torch import nn |
| 37 | |
| 38 | |
| 39 | # --------------------------------------------------------------------------- |
| 40 | # Config |
| 41 | # --------------------------------------------------------------------------- |
| 42 | @dataclass |
| 43 | class HyperConnectionConfig: |
| 44 | """Configuration shared by all HyperConnection variants.""" |
| 45 | |
| 46 | hc_count: int = 4 |
| 47 | hidden_size: int = 64 |
| 48 | params_dtype: torch.dtype = torch.bfloat16 |
| 49 | mtp_hc: bool = False |
| 50 | hc_lowrank: int = 16 |
| 51 | rms_norm_eps: float = 1e-6 |
| 52 | hc_per_branch_norm: bool = False |
| 53 | |
| 54 | |
| 55 | class GroupedGemmaRMSNorm(nn.Module): |
| 56 | def __init__( |
| 57 | self, |
| 58 | hidden_size: int, |
| 59 | eps: float, |
| 60 | group_size: int | None, |
| 61 | dtype: torch.dtype | None, |
| 62 | ) -> None: |
| 63 | super().__init__() |
| 64 | if group_size is not None and hidden_size % group_size: |
| 65 | raise ValueError( |
| 66 | f"hidden_size ({hidden_size}) must be divisible by " |
| 67 | f"group_size ({group_size})" |
| 68 | ) |
| 69 | self.variance_epsilon = eps |
| 70 | self.group_size = group_size |
| 71 | self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) |
| 72 | |
| 73 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 74 | input_dtype = hidden_states.dtype |
| 75 | hidden_states = hidden_states.float() |
| 76 | if self.group_size is None: |
| 77 | variance = hidden_states.square().mean(dim=-1, keepdim=True) |
| 78 | normalized = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
| 79 | else: |
| 80 | grouped = hidden_states.unflatten( |
| 81 | -1, (hidden_states.shape[-1] // self.group_size, self.group_size) |
| 82 | ) |
| 83 | variance = grouped.square().mean(dim=-1, keepdim=True) |
| 84 | normalized = ( |
| 85 | grouped * torch.rsqrt(variance + self.variance_epsilon) |
| 86 | ).flatten(-2) |
| 87 | return (normalized * (1.0 + self.weight.float())).to(input_dtype) |
| 88 | |
| 89 | |
| 90 | # --------------------------------------------------------------------------- |
| 91 | # Average-pooling variant |
| 92 | # --------------------------------------------------------------------------- |
| 93 | class HyperConnectionBase(nn.Module): |
| 94 | """Average-pooling HyperConnection (``hyperconnection_average``). |
| 95 | |
| 96 | Splits the incoming ``[..., HC*HS]`` tensor (HC outer, HS inner) into |
| 97 | ``HC`` parallel streams, averages them for the block input, and |
| 98 | broadcasts the block output back to every stream. |
| 99 | """ |
| 100 | |
| 101 | def __init__( |
| 102 | self, |
| 103 | config: HyperConnectionConfig, |
| 104 | layer_idx: int | None = None, |
| 105 | role: str | None = None, |
| 106 | ) -> None: |
| 107 | super().__init__() |
| 108 | self.config = config |
| 109 | self.hc_count = config.hc_count |
| 110 | self.hidden_size = config.hidden_size |
| 111 | self.layer_idx = layer_idx |
| 112 | self.role = role |
| 113 | |
| 114 | @property |
| 115 | def hyper_hidden_size(self) -> int: |
| 116 | return self.hc_count * self.hidden_size |
| 117 | |
| 118 | def mix(self, hyper_input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 119 | """Average the HC streams into a single block input.""" |
| 120 | assert hyper_input.shape[-1] == self.hc_count * self.hidden_size |
| 121 | # [*, HC, HS] — mean over HC (dim=-2). |
| 122 | unflat = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) |
| 123 | mixed_input = unflat.mean(dim=-2) |
| 124 | return mixed_input, hyper_input |
| 125 | |
| 126 | def combine( |
| 127 | self, block_output: torch.Tensor, residual: torch.Tensor |
| 128 | ) -> torch.Tensor: |
| 129 | """Broadcast the block output back to every stream.""" |
| 130 | assert residual.shape[-1] == self.hc_count * self.hidden_size |
| 131 | assert block_output.shape[-1] == self.hidden_size |
| 132 | residual_reshaped = residual.unflatten(-1, (self.hc_count, self.hidden_size)) |
| 133 | combined = residual_reshaped + block_output.unsqueeze(-2) |
| 134 | return combined.flatten(-2) |
| 135 | |
| 136 | |
| 137 | # --------------------------------------------------------------------------- |
| 138 | # Gated-residual variant |
| 139 | # --------------------------------------------------------------------------- |
| 140 | class GatedResidual(HyperConnectionBase): |
| 141 | """Gated HyperConnection with learnable low-rank mixing and injection. |
| 142 | |
| 143 | ``mix()`` applies GemmaRMSNorm per HC stream and projects through a |
| 144 | low-rank sigmoid gate to produce a single block input. ``combine()`` |
| 145 | injects the block output back into each stream through a learned |
| 146 | per-stream injection weight. |
| 147 | |
| 148 | This implementation uses only PyTorch operators. Tensor-parallel |
| 149 | collectives are supplied by its caller. |
| 150 | """ |
| 151 | |
| 152 | def __init__( |
| 153 | self, |
| 154 | config: HyperConnectionConfig, |
| 155 | layer_idx: int | None = None, |
| 156 | role: str | None = None, |
| 157 | use_mix: bool = True, |
| 158 | use_combine: bool = True, |
| 159 | ) -> None: |
| 160 | super().__init__(config, layer_idx, role) |
| 161 | norm_size = ( |
| 162 | self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size |
| 163 | ) |
| 164 | group_size = config.hidden_size if config.hc_per_branch_norm else None |
| 165 | # Normalize each H-sized HC stream independently while retaining a |
| 166 | # separate affine weight for every element of the HC*H layout. |
| 167 | self.hc_norm = GroupedGemmaRMSNorm( |
| 168 | norm_size, |
| 169 | eps=config.rms_norm_eps, |
| 170 | group_size=group_size, |
| 171 | dtype=config.params_dtype, |
| 172 | ) |
| 173 | |
| 174 | # -- raw Linear weights (checkpoint-compatible) ---------------------- |
| 175 | if use_mix: |
| 176 | self.input_mix_weight_down = nn.Linear( |
| 177 | self.hyper_hidden_size, |
| 178 | config.hc_lowrank, |
| 179 | bias=False, |
| 180 | dtype=config.params_dtype, |
| 181 | ) |
| 182 | self.input_mix_weight_up = nn.Linear( |
| 183 | config.hc_lowrank, |
| 184 | self.hyper_hidden_size, |
| 185 | bias=False, |
| 186 | dtype=config.params_dtype, |
| 187 | ) |
| 188 | if use_combine: |
| 189 | self.block_inject_weight = nn.Linear( |
| 190 | self.hyper_hidden_size, |
| 191 | self.hc_count, |
| 192 | bias=False, |
| 193 | dtype=config.params_dtype, |
| 194 | ) |
| 195 | |
| 196 | def _normalize(self, hyper_input: torch.Tensor) -> torch.Tensor: |
| 197 | if self.config.hc_per_branch_norm: |
| 198 | return self.hc_norm(hyper_input) |
| 199 | return self.hc_norm( |
| 200 | hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) |
| 201 | ).flatten(-2) |
| 202 | |
| 203 | def mix( |
| 204 | self, hyper_input: torch.Tensor |
| 205 | ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: |
| 206 | """Mix: RMSNorm -> low-rank gate -> gated mean.""" |
| 207 | assert hyper_input.shape[-1] == self.hc_count * self.hidden_size |
| 208 | if not hasattr(self, "input_mix_weight_down"): |
| 209 | raise RuntimeError("mix was disabled for this hyper-connection") |
| 210 | hyper_input_normed = self._normalize(hyper_input) |
| 211 | # Gate — original mix order: linear+silu then linear+sigmoid. |
| 212 | gate = F.silu( |
| 213 | F.linear(hyper_input_normed, self.input_mix_weight_down.weight) |
| 214 | / self.hc_count |
| 215 | ) |
| 216 | gate = torch.sigmoid(F.linear(gate, self.input_mix_weight_up.weight)).unflatten( |
| 217 | -1, (self.hc_count, self.hidden_size) |
| 218 | ) |
| 219 | mixed_input = ( |
| 220 | gate * hyper_input_normed.unflatten(-1, (self.hc_count, self.hidden_size)) |
| 221 | ).mean(dim=-2) |
| 222 | return mixed_input.to(hyper_input.dtype), (hyper_input, hyper_input_normed) |
| 223 | |
| 224 | def combine( |
| 225 | self, |
| 226 | block_output: torch.Tensor, |
| 227 | residuals: tuple[torch.Tensor, torch.Tensor], |
| 228 | ) -> torch.Tensor: |
| 229 | if not hasattr(self, "block_inject_weight"): |
| 230 | raise RuntimeError("combine was disabled for this hyper-connection") |
| 231 | hyper_input, hyper_input_normed = residuals |
| 232 | assert hyper_input.shape[-1] == self.hc_count * self.hidden_size |
| 233 | assert block_output.shape[-1] == self.hidden_size |
| 234 | residual = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) |
| 235 | # The paired mix keeps its normalized hyper input so combine uses the |
| 236 | # same HC module's injection weight. |
| 237 | injection_weight = 2.0 * torch.sigmoid( |
| 238 | F.linear(hyper_input_normed, self.block_inject_weight.weight) |
| 239 | / self.hc_count |
| 240 | ) |
| 241 | output = residual + block_output.unsqueeze(-2) * injection_weight.unsqueeze(-1) |
| 242 | return output.flatten(-2).to(hyper_input.dtype) |
| 243 | |
| 244 | |
| 245 | __all__ = [ |
| 246 | "GatedResidual", |
| 247 | "GroupedGemmaRMSNorm", |
| 248 | "HyperConnectionBase", |
| 249 | "HyperConnectionConfig", |
| 250 | ] |