Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / grouped_experts.py: 65%
226 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""NPU grouped routed-expert module."""
17from __future__ import annotations
19from collections.abc import Mapping
20from functools import partial
21from typing import Any
23import torch # pylint: disable=forbidden-backend-import
24from torch import nn # pylint: disable=forbidden-backend-import
25from transformers.activations import ACT2FN
26from hyper_parallel.components.checkpoint.weight_conversion import (
27 Transpose,
28 WeightConverter,
29 WeightRenaming,
30)
32from hyper_parallel.models.replacement import module_replacement
33from hyper_parallel.components.functional import (
34 grouped_matmul,
35 moe_token_permute,
36 moe_token_unpermute,
37 swiglu,
38)
41@module_replacement
42class GroupedExperts(nn.Module):
43 """Transformers-compatible NPU grouped routed experts.
45 Holds gate_up_proj/up_proj, down_proj (and optional bias1, bias2) as nn.Parameters.
46 Computation logic is handled by _grouped_gemm_expert_forward.
48 Args:
49 module: Source Transformers experts module.
50 module_fqn: Fully qualified name supplied by the replacement framework.
51 context: Additional replacement context.
52 """
54 @staticmethod
55 def _source_parameters(module: nn.Module, source_fc1_name: str) -> tuple[nn.Parameter, nn.Parameter]:
56 """Validate and return source expert weight parameters."""
57 missing = [name for name in (source_fc1_name, "down_proj") if not hasattr(module, name)]
58 if missing:
59 raise TypeError(f"GroupedExperts source module is missing: {missing}")
60 source_gate_up = getattr(module, source_fc1_name)
61 if not isinstance(source_gate_up, nn.Parameter):
62 raise TypeError(f"GroupedExperts requires {source_fc1_name} to be an nn.Parameter")
63 if not isinstance(module.down_proj, nn.Parameter):
64 raise TypeError("GroupedExperts requires down_proj to be an nn.Parameter")
65 return source_gate_up, module.down_proj
67 @staticmethod
68 def _source_dimensions(
69 module: nn.Module,
70 config: Any,
71 source_gate_up: nn.Parameter,
72 source_down: nn.Parameter,
73 gated_linear_unit: bool,
74 ) -> tuple[int, int, int]:
75 """Infer expert count and projection dimensions from the source module."""
76 num_experts = getattr(module, "num_local_experts", getattr(module, "num_experts", None))
77 if num_experts is None and config is not None:
78 num_experts = getattr(config, "n_routed_experts", getattr(config, "num_experts", None))
79 hidden_size = getattr(module, "hidden_size", getattr(module, "hidden_dim", None))
80 if hidden_size is None and config is not None:
81 hidden_size = getattr(config, "hidden_size", None)
82 intermediate_size = getattr(module, "intermediate_size", getattr(module, "intermediate_dim", None))
83 if intermediate_size is None and config is not None:
84 intermediate_size = getattr(config, "moe_intermediate_size", getattr(config, "intermediate_size", None))
85 source_is_transposed = getattr(module, "is_transposed", None)
86 if not gated_linear_unit and source_gate_up.dim() == 3:
87 hidden_size = GroupedExperts._infer_ungated_hidden_size(
88 source_gate_up, source_down, source_is_transposed, intermediate_size, hidden_size
89 )
90 if num_experts is None or hidden_size is None or intermediate_size is None:
91 raise ValueError("GroupedExperts requires expert count, hidden size, and intermediate size")
92 return num_experts, hidden_size, intermediate_size
94 @staticmethod
95 def _infer_ungated_hidden_size(
96 source_gate_up: nn.Parameter,
97 source_down: nn.Parameter,
98 source_is_transposed: bool | None,
99 intermediate_size: int | None,
100 default: int | None,
101 ) -> int | None:
102 """Infer hidden size for ungated three-dimensional source weights."""
103 if source_is_transposed is False:
104 return source_gate_up.shape[-1]
105 if source_is_transposed is True:
106 return source_gate_up.shape[-2]
107 if intermediate_size is None:
108 return default
109 if source_gate_up.shape[-2] == intermediate_size and source_down.shape[-1] == intermediate_size:
110 return source_gate_up.shape[-1]
111 if source_gate_up.shape[-1] == intermediate_size and source_down.shape[-2] == intermediate_size:
112 return source_gate_up.shape[-2]
113 return default
115 @staticmethod
116 def _activation(module: nn.Module, config: Any, gated_linear_unit: bool) -> Any:
117 """Resolve the source-compatible expert activation."""
118 hidden_act = getattr(config, "hidden_act", None)
119 if hidden_act is None:
120 hidden_act = getattr(config, "hidden_activation", None)
121 if hidden_act is None:
122 hidden_act = getattr(config, "mlp_hidden_act", None)
123 activation_func = getattr(module, "act_fn", None)
124 if activation_func is None:
125 if hidden_act is None:
126 hidden_act = "silu"
127 activation_func = ACT2FN[hidden_act]
128 source_apply_gate = getattr(module, "_apply_gate", None)
129 source_apply_gate_func = getattr(source_apply_gate, "__func__", source_apply_gate)
130 if (
131 gated_linear_unit
132 and source_apply_gate_func is not None
133 and getattr(source_apply_gate_func, "__name__", "") != "_default_apply_gate"
134 ):
135 raise ValueError("GroupedExperts custom _apply_gate requires a dedicated replacement")
136 if not gated_linear_unit:
137 return activation_func
138 if hidden_act == "silu" and bool(getattr(config, "use_fused_swiglu", True)):
139 return partial(swiglu, dim=-1)
141 def glu(x: torch.Tensor) -> torch.Tensor:
142 """Apply the configured gated linear unit."""
143 gate, value = torch.chunk(x, 2, dim=-1)
144 return activation_func(gate) * value
146 return glu
148 @staticmethod
149 def _source_layout(
150 source_shapes: tuple[tuple[int, ...], tuple[int, ...]],
151 source_is_transposed: bool | None,
152 expert_shapes: tuple[tuple[int, ...], tuple[int, ...]],
153 transformers_shapes: tuple[tuple[int, ...], tuple[int, ...]],
154 npu_shapes: tuple[tuple[int, ...], tuple[int, ...]],
155 ) -> tuple[bool, bool]:
156 """Return whether experts are two-dimensional and need transposition."""
157 if source_shapes == expert_shapes:
158 return True, False
159 if source_is_transposed is False:
160 if source_shapes != transformers_shapes:
161 raise ValueError("GroupedExperts source weights do not match their declared layout")
162 return False, True
163 if source_is_transposed is True:
164 if source_shapes != npu_shapes:
165 raise ValueError("GroupedExperts source weights do not match their declared layout")
166 return False, False
167 if source_shapes == transformers_shapes:
168 return False, True
169 if source_shapes == npu_shapes:
170 return False, False
171 raise ValueError("GroupedExperts source weight shapes are incompatible with the configured dimensions")
173 def _initialize_weights(
174 self,
175 source_gate_up: nn.Parameter,
176 source_down: nn.Parameter,
177 npu_gate_up_shape: tuple[int, ...],
178 npu_down_shape: tuple[int, ...],
179 gated_linear_unit: bool,
180 ) -> None:
181 """Create or reuse expert parameters in the target NPU layout."""
182 if self._transpose_source_weights:
183 target_gate_up = nn.Parameter(
184 torch.empty(npu_gate_up_shape, device=source_gate_up.device, dtype=source_gate_up.dtype),
185 requires_grad=source_gate_up.requires_grad,
186 )
187 self.down_proj = nn.Parameter(
188 torch.empty(npu_down_shape, device=source_down.device, dtype=source_down.dtype),
189 requires_grad=source_down.requires_grad,
190 )
191 else:
192 target_gate_up = source_gate_up
193 self.down_proj = source_down
194 if gated_linear_unit:
195 self.gate_up_proj = target_gate_up
196 else:
197 self.up_proj = target_gate_up
199 def _initialize_biases(self, module: nn.Module) -> None:
200 """Validate and materialize optional expert bias parameters."""
201 self._source_bias_names = (
202 "bias1" if hasattr(module, "bias1") else f"{self._source_fc1_name}_bias",
203 "bias2" if hasattr(module, "bias2") else "down_proj_bias",
204 )
205 source_bias1 = getattr(module, self._source_bias_names[0], None)
206 source_bias2 = getattr(module, self._source_bias_names[1], None)
207 if (source_bias1 is None) != (source_bias2 is None):
208 raise ValueError("GroupedExperts source projections must use the same bias policy")
209 self.add_bias = source_bias1 is not None
210 if not self.add_bias:
211 return
212 if not isinstance(source_bias1, nn.Parameter) or not isinstance(source_bias2, nn.Parameter):
213 raise TypeError("GroupedExperts source biases must be nn.Parameter instances")
214 target_biases = []
215 for source_name, target_name, source_bias in zip(
216 self._source_bias_names, ("bias1", "bias2"), (source_bias1, source_bias2)
217 ):
218 if source_name == target_name:
219 target_biases.append(source_bias)
220 else:
221 target_biases.append(
222 nn.Parameter(torch.empty_like(source_bias), requires_grad=source_bias.requires_grad)
223 )
224 # astroid cannot infer the length of a list built by appends in a
225 # loop; the zip above always yields exactly two entries.
226 self.bias1, self.bias2 = target_biases # pylint: disable=W0632
228 def __init__(
229 self,
230 *,
231 module: nn.Module,
232 module_fqn: str = "",
233 context: Mapping[str, Any] | None = None,
234 ) -> None:
235 """Build the high-performance module from a Transformers experts module."""
236 super().__init__()
237 del module_fqn, context
238 config = getattr(module, "config", None)
239 self.config = config
240 self.initializer_range = getattr(config, "initializer_range", None)
241 if bool(getattr(module, "has_bias", False)):
242 raise ValueError("GroupedExperts does not support source modules with has_bias=True")
243 gated_linear_unit = bool(
244 getattr(
245 module,
246 "has_gate",
247 getattr(config, "gated_linear_unit", True),
248 )
249 )
250 self._source_fc1_name = "gate_up_proj" if gated_linear_unit else "up_proj"
251 source_gate_up, source_down = self._source_parameters(module, self._source_fc1_name)
252 num_experts, hidden_size, intermediate_size = self._source_dimensions(
253 module, config, source_gate_up, source_down, gated_linear_unit
254 )
255 source_down = module.down_proj
256 source_is_transposed = getattr(module, "is_transposed", None)
258 self.num_local_experts = num_experts
259 self.num_experts = num_experts
260 self.hidden_size = hidden_size
261 self.intermediate_size = intermediate_size
262 self.router_gating_in_fp32 = bool(
263 getattr(module, "router_gating_in_fp32", True)
264 )
265 self.activation_func = self._activation(module, config, gated_linear_unit)
267 # GroupGemm state (NPU-only)
268 self._group_list = None
269 self._tokens_per_expert_gmm = None
271 self._initialize_projection_layout(
272 source_gate_up, source_down, source_is_transposed, gated_linear_unit
273 )
274 self._initialize_biases(module)
276 self.has_gate = gated_linear_unit
277 self.has_bias = self.add_bias
278 self.is_concatenated = bool(getattr(module, "is_concatenated", True))
279 if self.has_gate and not self.is_concatenated:
280 raise ValueError("GroupedExperts requires concatenated gate/up expert weights")
281 self.is_transposed = True
282 # Keep HF native initialization order stable; this replacement module
283 # resets its packed expert parameters after materialization instead.
284 self._hp_reset_after_materialization = True
285 self.train(module.training)
287 def _initialize_projection_layout(
288 self, source_gate_up, source_down, source_is_transposed, gated_linear_unit,
289 ) -> None:
290 """Resolve expert weight layouts and initialize their target parameters."""
291 hidden_size = self.hidden_size
292 intermediate_size = self.intermediate_size
293 fc1_output_size = intermediate_size
294 if gated_linear_unit:
295 fc1_output_size *= 2
297 npu_gate_up_shape = (
298 self.num_local_experts,
299 hidden_size,
300 fc1_output_size,
301 )
302 npu_down_shape = (
303 self.num_local_experts,
304 intermediate_size,
305 hidden_size,
306 )
307 transformers_gate_up_shape = (
308 self.num_local_experts,
309 fc1_output_size,
310 hidden_size,
311 )
312 transformers_down_shape = (
313 self.num_local_experts,
314 hidden_size,
315 intermediate_size,
316 )
317 experts_2d_gate_up_shape = (
318 self.num_local_experts * hidden_size,
319 fc1_output_size,
320 )
321 experts_2d_down_shape = (
322 self.num_local_experts * intermediate_size,
323 hidden_size,
324 )
325 source_shapes = (tuple(source_gate_up.shape), tuple(source_down.shape))
326 self.use_2d_experts, self._transpose_source_weights = self._source_layout(
327 source_shapes,
328 source_is_transposed,
329 (experts_2d_gate_up_shape, experts_2d_down_shape),
330 (transformers_gate_up_shape, transformers_down_shape),
331 (npu_gate_up_shape, npu_down_shape),
332 )
333 self._initialize_weights(
334 source_gate_up, source_down, npu_gate_up_shape, npu_down_shape, gated_linear_unit
335 )
337 def reset_parameters(self) -> None:
338 """Initialize grouped weights with the source model's configured standard deviation."""
339 if self.initializer_range is None:
340 raise ValueError(
341 "GroupedExperts random initialization requires config.initializer_range"
342 )
343 target_fc1 = self.gate_up_proj if self.has_gate else self.up_proj
344 nn.init.normal_(target_fc1, mean=0.0, std=self.initializer_range)
345 nn.init.normal_(self.down_proj, mean=0.0, std=self.initializer_range)
346 if self.add_bias:
347 nn.init.zeros_(self.bias1)
348 nn.init.zeros_(self.bias2)
350 def make_transforms(self) -> list[WeightRenaming | WeightConverter]:
351 """Describe reversible source-checkpoint to high-performance conversion."""
352 transforms: list[WeightRenaming | WeightConverter] = []
353 if self._transpose_source_weights:
354 transforms.extend(
355 [
356 WeightConverter(
357 source_patterns=self._source_fc1_name,
358 target_patterns=self._source_fc1_name,
359 operations=[Transpose(dim0=-2, dim1=-1)],
360 ),
361 WeightConverter(
362 source_patterns="down_proj",
363 target_patterns="down_proj",
364 operations=[Transpose(dim0=-2, dim1=-1)],
365 ),
366 ]
367 )
368 if self.add_bias:
369 for source_name, target_name in zip(
370 self._source_bias_names,
371 ("bias1", "bias2"),
372 ):
373 if source_name == target_name:
374 continue
375 transforms.append(
376 WeightRenaming(
377 source_patterns=source_name,
378 target_patterns=target_name,
379 )
380 )
381 return transforms
383 def _grouped_gemm_expert_forward(
384 self, gate_up_proj, down_proj, permuted, tokens_per_expert, permuted_probs
385 ):
386 """GroupGemm expert computation (NPU-only)."""
387 self._tokens_per_expert_gmm = tokens_per_expert.to(device=permuted.device)
388 self._group_list = self._tokens_per_expert_gmm.cumsum(dim=0)
390 if self.use_2d_experts:
391 gate_up_proj = gate_up_proj.view(self.num_local_experts, self.hidden_size, -1)
393 # up
394 if permuted.nelement() != 0:
395 fc1_output = grouped_matmul( # pylint: disable=not-callable
396 permuted, gate_up_proj, bias=None, group_list=self._group_list, group_type=0, group_list_type=0,
397 )
398 if self.add_bias:
399 b1 = self.bias1.view(self.num_local_experts, -1)
400 fc1_output = fc1_output + torch.repeat_interleave(b1, self._tokens_per_expert_gmm, dim=0)
401 else:
402 gate_up_proj_2d = gate_up_proj.view(self.hidden_size, -1)
403 fc1_output = torch.matmul(permuted, gate_up_proj_2d)
405 # down
406 if not self.router_gating_in_fp32 and permuted_probs is not None:
407 fc1_output = (
408 self.activation_func(fc1_output)
409 * permuted_probs.reshape(*fc1_output.shape[:-1], 1)
410 )
411 else:
412 fc1_output = self.activation_func(fc1_output)
414 if self.use_2d_experts:
415 down_proj = down_proj.view(self.num_local_experts, -1, self.hidden_size)
417 if fc1_output.nelement() != 0:
418 fc2_output = grouped_matmul( # pylint: disable=not-callable
419 fc1_output, down_proj, bias=None, group_list=self._group_list, group_type=0, group_list_type=0,
420 )
421 if self.add_bias:
422 b2 = self.bias2.view(self.num_local_experts, -1)
423 fc2_output = fc2_output + torch.repeat_interleave(
424 b2, self._tokens_per_expert_gmm, dim=0,
425 )
426 else:
427 down_proj_2d = down_proj.view(-1, self.hidden_size)
428 fc2_output = torch.matmul(fc1_output, down_proj_2d)
430 return fc2_output
432 def forward_expert_major(
433 self,
434 x: torch.Tensor,
435 num_tokens_per_expert: torch.Tensor,
436 scores: torch.Tensor | None = None,
437 ) -> torch.Tensor:
438 """Run grouped experts on expert-major routed tokens.
440 Args:
441 x: Expert-major input with shape ``[tokens, hidden_size]``.
442 num_tokens_per_expert: Token count for every local expert.
443 scores: Optional routing weight for every routed token.
445 Returns:
446 Expert-major output with shape ``[tokens, hidden_size]``.
447 """
448 gate_up_proj = self.gate_up_proj if self.has_gate else self.up_proj
449 down_proj = self.down_proj
450 self.num_local_experts = num_tokens_per_expert.shape[0]
452 return self._grouped_gemm_expert_forward(
453 gate_up_proj, down_proj, x, num_tokens_per_expert, scores
454 )
456 def forward(
457 self,
458 hidden_states: torch.Tensor,
459 top_k_index: torch.Tensor,
460 top_k_weights: torch.Tensor,
461 ) -> torch.Tensor:
462 """Run grouped experts with the Transformers Experts interface."""
463 hidden_shape = hidden_states.shape
464 hidden_states_flat = hidden_states.view(-1, hidden_states.shape[-1])
465 permuted_tokens, sorted_indices = moe_token_permute( # pylint: disable=not-callable
466 hidden_states_flat,
467 top_k_index,
468 )
469 flatten_probs = top_k_weights.view(-1)
470 permuted_probs = flatten_probs.index_select(0, sorted_indices)
471 tokens_per_expert = torch.histc(
472 top_k_index.view(-1).float(),
473 bins=self.num_experts,
474 min=0,
475 max=self.num_experts,
476 ).long()
477 expert_outputs = self.forward_expert_major(
478 permuted_tokens,
479 tokens_per_expert,
480 permuted_probs,
481 )
482 output = moe_token_unpermute( # pylint: disable=not-callable
483 expert_outputs,
484 sorted_indices,
485 top_k_weights,
486 )
487 return output.view(hidden_shape)