Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / modules / grouped_experts.py: 79%
96 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"""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.config import (
27 GROUPED_NDIM,
28 GROUP_LIST_COUNTS,
29 MATRIX_NDIM,
30)
31from hyper_parallel.components.quantization.functional.base_gmm_func import (
32 GroupedLinear,
33 _GroupedLinearFunction,
34)
35from hyper_parallel.components.quantization.ops.npu_mxfp8 import LowPrecisionCapabilityError
36from hyper_parallel.components.quantization.ops.npu_w4a8 import (
37 W4A8CapabilityError,
38)
40_EXPERT_PARAMETER_NAMES = ("gate_up_proj", "down_proj")
41_GATE_UP_BRANCHES = 2
44class GroupedExperts(nn.Module):
45 """Canonical low-precision shell for DeepSeek-V3 grouped experts.
47 The bound grouped-linear compute may select MXFP8, HiFloat8, native W4A8,
48 or fake W4A8. The shell owns no quantizer; it holds only the
49 ``grouped_linear`` instance selected by ``strategy_factory``.
50 """
52 def __init__(
53 self,
54 num_experts: int,
55 hidden_dim: int,
56 intermediate_dim: int,
57 *,
58 grouped_linear: GroupedLinear,
59 fqn: str = "",
60 ) -> None:
61 """Create an unbound packed-expert module."""
63 self._validate_grouped_linear(grouped_linear)
64 super().__init__()
65 self.gate_up_proj = nn.Parameter(
66 torch.empty(
67 num_experts,
68 _GATE_UP_BRANCHES * intermediate_dim,
69 hidden_dim,
70 )
71 )
72 self.down_proj = nn.Parameter(
73 torch.empty(num_experts, hidden_dim, intermediate_dim)
74 )
75 self.act_fn = nn.SiLU()
76 self.num_experts = num_experts
77 self.hidden_dim = hidden_dim
78 self.intermediate_dim = intermediate_dim
79 self.fqn = fqn
80 self.grouped_linear = grouped_linear
81 self.activation_limit = None
83 @staticmethod
84 def _validate_grouped_linear(grouped_linear: GroupedLinear) -> None:
85 """Require a plain strategy so module registration stays unchanged."""
87 if isinstance(grouped_linear, nn.Module) or not isinstance(
88 grouped_linear,
89 GroupedLinear,
90 ):
91 raise TypeError(
92 "grouped_linear must be a non-Module GroupedLinear instance"
93 )
95 @classmethod
96 def from_module(
97 cls,
98 source: nn.Module,
99 *,
100 fqn: str,
101 grouped_linear: GroupedLinear,
102 ) -> "GroupedExperts":
103 """Create a no-allocation shell retaining the source registrations."""
105 cls._validate_grouped_linear(grouped_linear)
106 parameter_names = tuple(source._parameters) # pylint: disable=protected-access
107 if parameter_names != _EXPERT_PARAMETER_NAMES:
108 raise TypeError(
109 f"{fqn!r} must register packed expert parameters "
110 f"{_EXPERT_PARAMETER_NAMES}, but got {parameter_names}."
111 )
112 gate_up_proj = source._parameters["gate_up_proj"] # pylint: disable=protected-access
113 down_proj = source._parameters["down_proj"] # pylint: disable=protected-access
114 if gate_up_proj is None or down_proj is None:
115 raise ValueError(f"{fqn!r} packed expert parameters cannot be None.")
116 if gate_up_proj.ndim != GROUPED_NDIM or down_proj.ndim != GROUPED_NDIM:
117 raise ValueError(
118 f"{fqn!r} packed expert weights must be three-dimensional."
119 )
120 expected_down_shape = (
121 gate_up_proj.shape[0],
122 gate_up_proj.shape[2],
123 gate_up_proj.shape[1] // _GATE_UP_BRANCHES,
124 )
125 if (
126 gate_up_proj.shape[1] % _GATE_UP_BRANCHES
127 or tuple(down_proj.shape) != expected_down_shape
128 ):
129 raise ValueError(
130 f"{fqn!r} has incompatible gate/up and down expert shapes: "
131 f"gate_up={tuple(gate_up_proj.shape)}, down={tuple(down_proj.shape)}."
132 )
134 # Create a no-allocation shell: cls.__new__(cls) resolves to
135 # object.__new__ (a staticmethod), so passing cls explicitly is correct;
136 # pylint mis-infers the binding of __new__ on a classmethod's cls.
137 converted = cls.__new__(cls) # pylint: disable=no-value-for-parameter
138 # Deliberately invoke nn.Module.__init__ directly so the shell does not
139 # re-run this class's allocating __init__.
140 nn.Module.__init__(converted) # pylint: disable=unnecessary-dunder-call
141 for name, parameter in source._parameters.items(): # pylint: disable=protected-access
142 converted.register_parameter(name, parameter)
143 for name, buffer in source._buffers.items(): # pylint: disable=protected-access
144 converted.register_buffer(
145 name,
146 buffer,
147 persistent=name not in source._non_persistent_buffers_set, # pylint: disable=protected-access
148 )
149 for name, module in source._modules.items(): # pylint: disable=protected-access
150 converted.add_module(name, module)
151 converted.act_fn = getattr(source, "act_fn", nn.SiLU())
153 converted.num_experts = gate_up_proj.shape[0]
154 converted.hidden_dim = gate_up_proj.shape[2]
155 converted.intermediate_dim = gate_up_proj.shape[1] // _GATE_UP_BRANCHES
156 converted.fqn = fqn
157 converted.grouped_linear = grouped_linear
158 converted.activation_limit = getattr(source, "limit", None)
159 if hasattr(source, "config"):
160 converted.config = source.config
161 converted.training = source.training
162 return converted
164 def _grouped_forward(
165 self,
166 sorted_inputs: torch.Tensor,
167 tokens_per_expert: torch.Tensor,
168 ) -> torch.Tensor:
169 """Run already expert-major tokens through the two low-precision GMMs."""
171 try:
172 gate_up = _GroupedLinearFunction.apply(
173 sorted_inputs,
174 self.gate_up_proj,
175 tokens_per_expert,
176 self.grouped_linear,
177 GROUP_LIST_COUNTS,
178 )
179 gate, up = gate_up.chunk(_GATE_UP_BRANCHES, dim=-1)
180 if self.activation_limit is None:
181 intermediate = self.act_fn(gate) * up
182 else:
183 gate = gate.clamp(max=self.activation_limit)
184 up = up.clamp(min=-self.activation_limit, max=self.activation_limit)
185 intermediate = self.act_fn(gate) * up
186 return _GroupedLinearFunction.apply(
187 intermediate,
188 self.down_proj,
189 tokens_per_expert,
190 self.grouped_linear,
191 GROUP_LIST_COUNTS,
192 )
193 except (LowPrecisionCapabilityError, W4A8CapabilityError) as exc:
194 target = self.fqn or "<unknown>"
195 raise LowPrecisionCapabilityError(
196 f"Low-precision grouped-expert target {target!r} cannot run: {exc}"
197 ) from exc
199 def forward(
200 self,
201 hidden_states: torch.Tensor,
202 top_k_index: torch.Tensor,
203 top_k_weights: torch.Tensor,
204 ) -> torch.Tensor:
205 """Sort routed tokens, execute grouped experts, and restore token order."""
207 if hidden_states.ndim != MATRIX_NDIM:
208 raise ValueError(
209 "Grouped experts require two-dimensional hidden_states, "
210 f"but got shape {tuple(hidden_states.shape)}."
211 )
212 if (
213 top_k_index.shape != top_k_weights.shape
214 or top_k_index.ndim != MATRIX_NDIM
215 ):
216 raise ValueError(
217 "top_k_index and top_k_weights must have the same two-dimensional shape."
218 )
219 if top_k_index.shape[0] != hidden_states.shape[0]:
220 raise ValueError(
221 "The routed token count must match hidden_states: "
222 f"routes={top_k_index.shape[0]}, tokens={hidden_states.shape[0]}."
223 )
225 input_dtype = hidden_states.dtype
226 token_count = hidden_states.shape[0]
227 top_k = top_k_index.shape[1]
228 source_token_indices = torch.arange(
229 token_count,
230 device=hidden_states.device,
231 ).unsqueeze(1).expand(-1, top_k).reshape(-1)
232 flattened_expert_indices = top_k_index.reshape(-1)
233 expert_order = flattened_expert_indices.argsort()
234 sorted_inputs = hidden_states[source_token_indices[expert_order]]
235 tokens_per_expert = torch.bincount(
236 flattened_expert_indices,
237 minlength=self.num_experts,
238 )
239 # bincount rejects negative indices; an extra bin identifies an upper bound.
240 if tokens_per_expert.numel() != self.num_experts:
241 raise ValueError(
242 f"Expert indices must be in [0, {self.num_experts})."
243 )
244 # bincount returns per-expert counts, so the incoming group-list type
245 # is 1. Each grouped-linear strategy adapts this only if its backend
246 # requires cumulative offsets instead.
247 sorted_outputs = self._grouped_forward(sorted_inputs, tokens_per_expert)
248 sorted_weights = top_k_weights.reshape(-1)[expert_order]
249 sorted_outputs = sorted_outputs * sorted_weights.unsqueeze(-1)
251 inverse_order = torch.empty_like(expert_order)
252 inverse_order[expert_order] = torch.arange(
253 expert_order.shape[0],
254 device=expert_order.device,
255 )
256 # Accumulate every routed contribution before rounding once at the
257 # module boundary. GroupedExperts preserves the input activation
258 # dtype even when FP32 routing weights promote the weighted outputs.
259 return sorted_outputs[inverse_order].view(
260 token_count,
261 top_k,
262 self.hidden_dim,
263 ).sum(dim=1).to(dtype=input_dtype)
266__all__ = ["GroupedExperts"]