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