Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / modules / grouped_experts.py: 78%
88 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"""Canonical low-precision adapter for DeepSeek-V3 grouped experts.
17``GroupedExperts`` is the single canonical grouped-experts shell: it only
18orchestrates routing, activation, and token ordering. Gate/up and down
19projections run through the bound ``GroupedLinear``, whose
20instance is picked exclusively in ``strategy_factory``.
21"""
23import torch # pylint: disable=forbidden-backend-import
24from torch import nn # pylint: disable=forbidden-backend-import
26from hyper_parallel.components.quantization.functional.base_gmm_func import (
27 GroupedLinear,
28 _GroupedLinearFunction,
29)
30from hyper_parallel.components.quantization.ops.npu_mxfp8 import LowPrecisionCapabilityError
31from hyper_parallel.components.quantization.ops.npu_w4a8 import (
32 W4A8CapabilityError,
33)
35_EXPERT_PARAMETER_NAMES = ("gate_up_proj", "down_proj")
38class GroupedExperts(nn.Module):
39 """Canonical low-precision shell for DeepSeek-V3 grouped experts.
41 The bound grouped-linear compute may select MXFP8, HiFloat8, native W4A8,
42 or fake W4A8. The shell owns no quantizer; it holds only the
43 ``grouped_linear`` instance selected by ``strategy_factory``.
44 """
46 def __init__(
47 self,
48 num_experts: int,
49 hidden_dim: int,
50 intermediate_dim: int,
51 *,
52 grouped_linear: GroupedLinear,
53 fqn: str = "",
54 ) -> None:
55 """Create an unbound packed-expert module."""
57 self._validate_grouped_linear(grouped_linear)
58 super().__init__()
59 self.gate_up_proj = nn.Parameter(
60 torch.empty(num_experts, 2 * intermediate_dim, hidden_dim)
61 )
62 self.down_proj = nn.Parameter(
63 torch.empty(num_experts, hidden_dim, intermediate_dim)
64 )
65 self.act_fn = nn.SiLU()
66 self.num_experts = num_experts
67 self.hidden_dim = hidden_dim
68 self.intermediate_dim = intermediate_dim
69 self.fqn = fqn
70 self.grouped_linear = grouped_linear
72 @staticmethod
73 def _validate_grouped_linear(grouped_linear: GroupedLinear) -> None:
74 """Require a plain strategy so module registration stays unchanged."""
76 if isinstance(grouped_linear, nn.Module) or not isinstance(
77 grouped_linear,
78 GroupedLinear,
79 ):
80 raise TypeError(
81 "grouped_linear must be a non-Module GroupedLinear instance"
82 )
84 @classmethod
85 def from_module(
86 cls,
87 source: nn.Module,
88 *,
89 fqn: str,
90 grouped_linear: GroupedLinear,
91 ) -> "GroupedExperts":
92 """Create a no-allocation shell retaining the source registrations."""
94 cls._validate_grouped_linear(grouped_linear)
95 parameter_names = tuple(source._parameters) # pylint: disable=protected-access
96 if parameter_names != _EXPERT_PARAMETER_NAMES:
97 raise TypeError(
98 f"{fqn!r} must register packed expert parameters "
99 f"{_EXPERT_PARAMETER_NAMES}, but got {parameter_names}."
100 )
101 gate_up_proj = source._parameters["gate_up_proj"] # pylint: disable=protected-access
102 down_proj = source._parameters["down_proj"] # pylint: disable=protected-access
103 if gate_up_proj is None or down_proj is None:
104 raise ValueError(f"{fqn!r} packed expert parameters cannot be None.")
105 if gate_up_proj.ndim != 3 or down_proj.ndim != 3:
106 raise ValueError(
107 f"{fqn!r} packed expert weights must be three-dimensional."
108 )
109 expected_down_shape = (
110 gate_up_proj.shape[0],
111 gate_up_proj.shape[2],
112 gate_up_proj.shape[1] // 2,
113 )
114 if gate_up_proj.shape[1] % 2 or tuple(down_proj.shape) != expected_down_shape:
115 raise ValueError(
116 f"{fqn!r} has incompatible gate/up and down expert shapes: "
117 f"gate_up={tuple(gate_up_proj.shape)}, down={tuple(down_proj.shape)}."
118 )
120 # Create a no-allocation shell: cls.__new__(cls) resolves to
121 # object.__new__ (a staticmethod), so passing cls explicitly is correct;
122 # pylint mis-infers the binding of __new__ on a classmethod's cls.
123 converted = cls.__new__(cls) # pylint: disable=no-value-for-parameter
124 # Deliberately invoke nn.Module.__init__ directly so the shell does not
125 # re-run this class's allocating __init__.
126 nn.Module.__init__(converted) # pylint: disable=unnecessary-dunder-call
127 for name, parameter in source._parameters.items(): # pylint: disable=protected-access
128 converted.register_parameter(name, parameter)
129 for name, buffer in source._buffers.items(): # pylint: disable=protected-access
130 converted.register_buffer(
131 name,
132 buffer,
133 persistent=name not in source._non_persistent_buffers_set, # pylint: disable=protected-access
134 )
135 for name, module in source._modules.items(): # pylint: disable=protected-access
136 converted.add_module(name, module)
138 converted.num_experts = gate_up_proj.shape[0]
139 converted.hidden_dim = gate_up_proj.shape[2]
140 converted.intermediate_dim = gate_up_proj.shape[1] // 2
141 converted.fqn = fqn
142 converted.grouped_linear = grouped_linear
143 if hasattr(source, "config"):
144 converted.config = source.config
145 converted.training = source.training
146 return converted
148 def _grouped_forward(
149 self,
150 sorted_inputs: torch.Tensor,
151 tokens_per_expert: torch.Tensor,
152 ) -> torch.Tensor:
153 """Run already expert-major tokens through the two low-precision GMMs."""
155 try:
156 gate_up = _GroupedLinearFunction.apply(
157 sorted_inputs,
158 self.gate_up_proj,
159 tokens_per_expert,
160 self.grouped_linear,
161 1,
162 )
163 gate, up = gate_up.chunk(2, dim=-1)
164 intermediate = self.act_fn(gate) * up
165 return _GroupedLinearFunction.apply(
166 intermediate,
167 self.down_proj,
168 tokens_per_expert,
169 self.grouped_linear,
170 1,
171 )
172 except (LowPrecisionCapabilityError, W4A8CapabilityError) as exc:
173 target = self.fqn or "<unknown>"
174 raise LowPrecisionCapabilityError(
175 f"Low-precision grouped-expert target {target!r} cannot run: {exc}"
176 ) from exc
178 def forward(
179 self,
180 hidden_states: torch.Tensor,
181 top_k_index: torch.Tensor,
182 top_k_weights: torch.Tensor,
183 ) -> torch.Tensor:
184 """Sort routed tokens, execute grouped experts, and restore token order."""
186 if hidden_states.ndim != 2:
187 raise ValueError(
188 "Grouped experts require two-dimensional hidden_states, "
189 f"but got shape {tuple(hidden_states.shape)}."
190 )
191 if top_k_index.shape != top_k_weights.shape or top_k_index.ndim != 2:
192 raise ValueError(
193 "top_k_index and top_k_weights must have the same two-dimensional shape."
194 )
195 if top_k_index.shape[0] != hidden_states.shape[0]:
196 raise ValueError(
197 "The routed token count must match hidden_states: "
198 f"routes={top_k_index.shape[0]}, tokens={hidden_states.shape[0]}."
199 )
201 token_count = hidden_states.shape[0]
202 top_k = top_k_index.shape[1]
203 source_token_indices = torch.arange(
204 token_count,
205 device=hidden_states.device,
206 ).unsqueeze(1).expand(-1, top_k).reshape(-1)
207 flattened_expert_indices = top_k_index.reshape(-1)
208 expert_order = flattened_expert_indices.argsort()
209 sorted_inputs = hidden_states[source_token_indices[expert_order]]
210 tokens_per_expert = torch.bincount(
211 flattened_expert_indices,
212 minlength=self.num_experts,
213 )
214 # bincount rejects negative indices; an extra bin identifies an upper bound.
215 if tokens_per_expert.numel() != self.num_experts:
216 raise ValueError(
217 f"Expert indices must be in [0, {self.num_experts})."
218 )
219 sorted_outputs = self._grouped_forward(sorted_inputs, tokens_per_expert)
220 sorted_weights = top_k_weights.reshape(-1)[expert_order]
221 sorted_outputs = sorted_outputs * sorted_weights.unsqueeze(-1)
222 # Only HiFloat8 historically rounded each routed contribution before
223 # reduction; the MX strategies kept promotion from routing weights.
224 if self.grouped_linear.CAST_ROUTED_OUTPUT_TO_INPUT_DTYPE:
225 sorted_outputs = sorted_outputs.to(hidden_states.dtype)
227 inverse_order = torch.empty_like(expert_order)
228 inverse_order[expert_order] = torch.arange(
229 expert_order.shape[0],
230 device=expert_order.device,
231 )
232 return sorted_outputs[inverse_order].view(
233 token_count,
234 top_k,
235 self.hidden_dim,
236 ).sum(dim=1)
239__all__ = ["GroupedExperts"]