Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_gate / module.py: 84%
112 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Model-facing Torch MegaGate router module."""
17from __future__ import annotations
19__all__ = ["MegaGate"]
21import math
22from collections.abc import Callable
23from functools import lru_cache
24from typing import Any
26import torch
27from torch.nn import functional
29from hyper_multicore.modules._precision import preserve_fp32_parameters
30from hyper_multicore.modules.mega_gate.function import mega_gate
31from hyper_multicore.modules.mega_gate.plan import MegaGatePlan, build_mega_gate_plan
34def _reset_mega_gate_parameters(
35 weight: torch.nn.Parameter,
36 bias: torch.nn.Parameter,
37 bias_vl: torch.nn.Parameter | None,
38 initializer_range: float,
39) -> None:
40 """Initialize one MegaGate parameter set."""
41 torch.nn.init.normal_(weight, mean=0.0, std=initializer_range)
42 torch.nn.init.zeros_(bias)
43 if bias_vl is not None:
44 torch.nn.init.zeros_(bias_vl)
47def _create_mega_gate_parameters(
48 hidden_size: int,
49 num_experts: int,
50 vision_enabled: bool,
51 initializer_range: float,
52) -> tuple[torch.nn.Parameter, torch.nn.Parameter, torch.nn.Parameter | None]:
53 """Create independently initialized projection and correction biases."""
54 weight = torch.nn.Parameter(torch.empty(num_experts, hidden_size))
55 bias = torch.nn.Parameter(torch.empty(num_experts, dtype=torch.float32))
56 bias_vl = (
57 torch.nn.Parameter(torch.empty(num_experts, dtype=torch.float32))
58 if vision_enabled
59 else None
60 )
61 _reset_mega_gate_parameters(weight, bias, bias_vl, initializer_range)
62 return weight, bias, bias_vl
65@lru_cache(maxsize=None)
66def _plan_for_device(device: torch.device) -> MegaGatePlan:
67 """Build one immutable MegaGate plan per process and NPU device."""
68 return build_mega_gate_plan(device)
71class MegaGate(torch.nn.Module):
72 """Project hidden states and select experts with MegaGate routing.
74 The module owns its projection and correction-bias parameters. NPU
75 sqrt-softplus calls use HyperMegaGate Route; other scoring functions and
76 non-NPU tensors use the equivalent Torch implementation.
77 """
79 def __init__(
80 self,
81 *,
82 hidden_size: int,
83 num_experts: int,
84 top_k: int,
85 scoring_func: str,
86 routed_scaling_factor: float,
87 vision_enabled: bool = False,
88 initializer_range: float = 0.02,
89 ) -> None:
90 """Initialize a standalone MegaGate router.
92 Args:
93 hidden_size: Size of the final hidden-state dimension.
94 num_experts: Number of routed experts.
95 top_k: Number of experts selected for each token.
96 scoring_func: ``sqrtsoftplus``, ``softmax`` or ``sigmoid``.
97 routed_scaling_factor: Finite multiplier applied to routing weights.
98 vision_enabled: Create a separate correction bias for visual tokens.
99 initializer_range: Standard deviation of the normal weight initialization.
101 Raises:
102 ValueError: An extent, scoring function, scale or initializer is invalid.
103 """
104 super().__init__()
105 if hidden_size <= 0 or num_experts <= 0:
106 raise ValueError("hidden_size and num_experts must be positive")
107 if top_k <= 0 or top_k > num_experts:
108 raise ValueError("top_k must be in [1, num_experts]")
109 if scoring_func not in ("sqrtsoftplus", "softmax", "sigmoid"):
110 raise ValueError(f"unsupported scoring_func: {scoring_func!r}")
111 if not math.isfinite(routed_scaling_factor):
112 raise ValueError("routed_scaling_factor must be finite")
113 if not math.isfinite(initializer_range) or initializer_range < 0.0:
114 raise ValueError("initializer_range must be finite and nonnegative")
115 self.hidden_size = int(hidden_size)
116 self.num_experts = int(num_experts)
117 self.top_k = int(top_k)
118 self.scoring_func = scoring_func
119 self.routed_scaling_factor = float(routed_scaling_factor)
120 self.initializer_range = float(initializer_range)
121 self.weight, self.bias, bias_vl = _create_mega_gate_parameters(
122 self.hidden_size,
123 self.num_experts,
124 vision_enabled,
125 self.initializer_range,
126 )
127 if bias_vl is None:
128 self.register_parameter("bias_vl", None)
129 else:
130 self.bias_vl = bias_vl
132 def reset_parameters(self) -> None:
133 """Reset projection and correction-bias parameters."""
134 _reset_mega_gate_parameters(
135 self.weight,
136 self.bias,
137 self.bias_vl,
138 self.initializer_range,
139 )
141 @classmethod
142 def from_module(cls, module: torch.nn.Module, **configuration: Any) -> MegaGate:
143 """Bind an existing router's projection and correction-bias Parameters.
145 Args:
146 module: Router with weight [E, H], optional correction bias [E] and bias_vl [E].
147 Its bias is a TopK correction, not a linear-projection bias.
148 **configuration: top_k, scoring_func and routed_scaling_factor, when not present on module.
149 vision_enabled defaults to whether the source provides bias_vl; it cannot disable an existing bias_vl.
151 Returns:
152 Router sharing the original Parameters. Construct the optimizer after module replacement.
153 """
154 weight = getattr(module, "weight", None)
155 if not isinstance(weight, torch.nn.Parameter) or weight.ndim != 2:
156 raise ValueError("source router requires a weight Parameter shaped [num_experts, hidden_size]")
157 options = {}
158 for name in ("top_k", "scoring_func", "routed_scaling_factor"):
159 value = configuration.pop(name, getattr(module, name, None))
160 if value is None:
161 raise ValueError(f"Provide {name} or expose it on the source router")
162 options[name] = value
163 bias_vl = getattr(module, "bias_vl", None)
164 vision_enabled = configuration.pop("vision_enabled", bias_vl is not None)
165 if bias_vl is not None and not vision_enabled:
166 raise ValueError("source bias_vl requires vision_enabled=True")
167 result = cls(hidden_size=weight.shape[1], num_experts=weight.shape[0],
168 top_k=options["top_k"], scoring_func=options["scoring_func"],
169 routed_scaling_factor=options["routed_scaling_factor"],
170 vision_enabled=vision_enabled, **configuration).to(device=weight.device)
171 result.weight = weight
172 for name in ("bias", "bias_vl"):
173 value = getattr(module, name, None)
174 if value is not None:
175 if (not isinstance(value, torch.nn.Parameter) or value.shape != (weight.shape[0],)
176 or value.dtype != torch.float32 or value.device != weight.device):
177 raise ValueError(f"source {name} must be an FP32 Parameter [num_experts] on the weight device")
178 setattr(result, name, value)
179 return result
181 def _apply(self, fn: Callable, recurse: bool = True) -> MegaGate:
182 """Keep correction biases FP32 under enclosing model dtype conversion."""
183 return super()._apply(preserve_fp32_parameters(self, fn, ("bias", "bias_vl")), recurse=recurse)
185 def _plan(self, hidden_states: torch.Tensor) -> MegaGatePlan:
186 """Return the process-level plan for an NPU tensor."""
187 return _plan_for_device(hidden_states.device)
189 def _torch_forward(
190 self,
191 hidden_states: torch.Tensor,
192 image_mask: torch.Tensor | None,
193 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
194 """Run the DeepSeek V4.1 reference operation order with Torch ops."""
195 flattened = hidden_states.reshape(-1, self.hidden_size)
196 logits = functional.linear( # pylint: disable=not-callable
197 flattened.float(), self.weight.float()
198 )
199 if self.scoring_func == "sqrtsoftplus":
200 scores = functional.softplus(logits).sqrt() # pylint: disable=not-callable
201 elif self.scoring_func == "softmax":
202 scores = logits.softmax(dim=-1)
203 elif self.scoring_func == "sigmoid":
204 scores = logits.sigmoid()
205 else:
206 raise ValueError(f"Unsupported V4.1 router scoring function: {self.scoring_func!r}")
208 correction_bias = self.bias
209 if image_mask is not None:
210 if self.bias_vl is not None:
211 correction_bias = torch.where(
212 image_mask.reshape(-1, 1),
213 self.bias_vl.unsqueeze(0),
214 self.bias.unsqueeze(0),
215 )
216 expert_indices = torch.topk(
217 scores + correction_bias,
218 self.top_k,
219 dim=-1,
220 sorted=False,
221 ).indices
222 routing_weights = scores.gather(1, expert_indices)
223 if self.top_k > 1:
224 routing_weights = routing_weights / (
225 routing_weights.sum(dim=-1, keepdim=True) + 1.0e-20
226 )
227 return logits, routing_weights * self.routed_scaling_factor, expert_indices
229 def forward(
230 self,
231 hidden_states: torch.Tensor,
232 image_mask: torch.Tensor | None = None,
233 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
234 """Project hidden states and select experts for every flattened token.
236 Args:
237 hidden_states: Tensor shaped ``[batch, sequence, hidden_size]``.
238 The projection follows the active autocast policy. The
239 logits are converted to FP32 before the native NPU Route.
240 image_mask: Optional BOOL tensor shaped ``[batch, sequence]``.
241 When visual bias is enabled, true entries select ``bias_vl``
242 and false entries select ``bias``. Native visual routing
243 requires a contiguous mask on the same NPU as
244 ``hidden_states``. Passing a mask without vision_enabled=True raises ValueError.
246 Returns:
247 A tuple containing logits ``[tokens, num_experts]`` in the
248 projection output dtype, FP32 routing weights
249 ``[tokens, top_k]`` and INT64 expert indices ``[tokens, top_k]``.
251 Raises:
252 ValueError: ``image_mask`` does not have shape
253 ``[batch, sequence]`` or a native visual-routing mask is not
254 contiguous.
255 RuntimeError: Native loading, dtype/device validation or execution
256 fails.
258 Note:
259 The native NPU Route supports first-order differentiation only.
260 Differentiating through RouteGrad is unsupported.
261 """
262 if image_mask is not None and self.bias_vl is None:
263 raise ValueError("image_mask requires vision_enabled=True when constructing MegaGate")
264 if image_mask is not None and image_mask.shape != hidden_states.shape[:2]:
265 raise ValueError("image_mask must have shape [batch, sequence]")
266 if (
267 self.scoring_func != "sqrtsoftplus"
268 or hidden_states.device.type != "npu"
269 or hidden_states.numel() == 0
270 ):
271 return self._torch_forward(hidden_states, image_mask)
273 vision_bias = self.bias_vl
274 use_vision_bias = image_mask is not None and vision_bias is not None
275 flat_image_mask = None
276 if use_vision_bias:
277 if not image_mask.is_contiguous():
278 raise ValueError("image_mask must be contiguous for native MegaGate")
279 flat_image_mask = image_mask.view(-1)
280 return mega_gate(
281 hidden_states,
282 self.weight,
283 self.bias,
284 self._plan(hidden_states),
285 vision_bias=vision_bias if use_vision_bias else None,
286 image_mask=flat_image_mask,
287 top_k=self.top_k,
288 routed_scaling_factor=self.routed_scaling_factor,
289 )