Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / base_gmm_func.py: 96%
106 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"""Shared grouped-linear autograd flow for native and fake-QAT strategies.
17The lifecycle prepares forward operands, executes GMM, retains only the views
18required by backward, prepares ``grad_output``, and executes dgrad/wgrad.
19Format-specific strategies provide quantization, direction-retention, group
20normalization, and backend calls; the shared template contains no format or
21operator dispatch.
22"""
24from abc import ABC, abstractmethod
25from typing import Optional
27import torch # pylint: disable=forbidden-backend-import
29from hyper_parallel.components.quantization.tensor import (
30 QuantizedTensor,
31 QuantizedTensorStorage,
32)
33from hyper_parallel.components.quantization.functional._saved_quantized import (
34 restore_quantized_operands,
35 save_quantized_operands,
36)
39class GroupedLinear(ABC):
40 """Own the common grouped-linear flow and expose format-specific hooks.
42 The weight contract is fixed to the DeepSeek-V3 live layout
43 ``[E, O, K]`` (contracting ``K`` at ``-1``): the adapter passes the packed
44 source parameters through unchanged. ``forward`` transposes each weight
45 once, inside the flow, to the GMM-ready ``[E, K, N]`` that strategies and
46 operators consume; ``backward`` transposes the wgrad result back to
47 ``[E, O, K]``.
48 """
50 FORMAT_NAME: str
51 CAST_ROUTED_OUTPUT_TO_INPUT_DTYPE = False
53 def _validate_inputs(
54 self,
55 inputs: torch.Tensor,
56 weight: torch.Tensor,
57 group_list: torch.Tensor,
58 group_list_type: int,
59 ) -> None:
60 """Validate the common live-weight and group metadata contract."""
62 validate_grouped_linear_inputs(
63 inputs,
64 weight,
65 group_list,
66 group_list_type,
67 weight_input_dim=-1,
68 name=self.FORMAT_NAME,
69 )
71 @abstractmethod
72 def normalize_group_list(
73 self,
74 group_list: torch.Tensor,
75 group_list_type: int,
76 ) -> tuple[torch.Tensor, int]:
77 """Adapt grouped metadata to this backend's fixed contract.
79 The adapter may hand per-expert token counts (``group_list_type=1``)
80 or cumulative offsets (``group_list_type=0``). MXFP8 and HiFloat8
81 accept either representation; W4A8 strategies normalize counts to
82 cumulative offsets and leave existing offsets unchanged.
83 """
85 @abstractmethod
86 def output_features(self, weight: torch.Tensor) -> int:
87 """Return the output width ``O`` in the live ``[E, O, K]`` layout."""
89 @abstractmethod
90 def weight_quantization_directions(
91 self,
92 needs_grad_input: bool,
93 ) -> tuple[bool, bool]:
94 """Return the row/column weight views needed for forward and dgrad."""
96 @abstractmethod
97 def retain_weight_backward(
98 self,
99 weight_quant: QuantizedTensorStorage,
100 needs_grad_input: bool,
101 ) -> None:
102 """Release the forward-only weight view after the forward GMM.
104 Each strategy keeps the direction consumed by its dgrad operator.
105 Input and gradient views are released inline in the template because
106 their lifecycle does not vary by format.
107 """
109 @abstractmethod
110 def quantize_input(
111 self,
112 inputs: torch.Tensor,
113 *,
114 rowwise: bool,
115 colwise: bool,
116 group_list: torch.Tensor,
117 group_list_type: int,
118 ) -> QuantizedTensor:
119 """Create the format-specific activation representation."""
121 @abstractmethod
122 def quantize_weight(
123 self,
124 weight: torch.Tensor,
125 *,
126 rowwise: bool,
127 colwise: bool,
128 ) -> QuantizedTensor:
129 """Create the format-specific weight representation."""
131 @abstractmethod
132 def quantize_grad_output(
133 self,
134 grad_output: torch.Tensor,
135 *,
136 rowwise: bool,
137 colwise: bool,
138 group_list: torch.Tensor,
139 group_list_type: int,
140 ) -> QuantizedTensor:
141 """Create the format-specific gradient-output representation."""
143 @abstractmethod
144 def grouped_matmul(
145 self,
146 left: QuantizedTensor,
147 right: QuantizedTensor,
148 *,
149 layout: str,
150 group_list: torch.Tensor,
151 group_type: int,
152 group_list_type: int,
153 output_dtype: torch.dtype,
154 ) -> torch.Tensor:
155 """Execute one format-specific grouped matrix multiplication."""
157 def forward(
158 self,
159 ctx: torch.autograd.function.FunctionCtx,
160 inputs: torch.Tensor,
161 weight: torch.Tensor,
162 group_list: torch.Tensor,
163 group_list_type: int,
164 ) -> torch.Tensor:
165 """Run the complete forward lifecycle using this strategy's hooks."""
167 # 共享校验门禁——live [E,O,K] 几何/group 不变量，
168 # 错误前缀取 FORMAT_NAME
169 self._validate_inputs(
170 inputs,
171 weight,
172 group_list,
173 group_list_type,
174 )
175 # hook: group 元数据归一化——W4A8 覆写为 counts→offsets
176 effective_group_list, effective_group_list_type = self.normalize_group_list(
177 group_list,
178 group_list_type,
179 )
180 ctx.input_shape = inputs.shape
181 ctx.input_dtype = inputs.dtype
182 ctx.input_device = inputs.device
183 ctx.weight_shape = weight.shape
184 ctx.weight_dtype = weight.dtype
185 ctx.weight_device = weight.device
186 ctx.group_list_type = effective_group_list_type
187 ctx.strategy = self
188 ctx.empty_input = inputs.shape[0] == 0
189 if ctx.empty_input:
190 # abstract: 给出输出宽 O，空输入也能产出 (0,O)
191 ctx.save_for_backward(effective_group_list)
192 return inputs.new_empty((0, self.output_features(weight)))
194 # The adapter hands live ``[E, O, K]`` weights; transpose once to the
195 # GMM-ready ``[E, K, N]`` that quantizers and operators consume.
196 gmm_weight = weight.transpose(-2, -1).contiguous()
197 needs_grad_input = inputs.requires_grad
198 needs_grad_weight = weight.requires_grad
199 # abstract: 量化激活（行视图给 forward，列视图给 wgrad）
200 input_quant = self.quantize_input(
201 inputs,
202 rowwise=True,
203 colwise=needs_grad_weight,
204 group_list=effective_group_list,
205 group_list_type=effective_group_list_type,
206 )
207 # hook: 决定要构建的 weight 视图——MXFP8 列向 fwd、W4A8 行向 fwd
208 weight_rowwise, weight_colwise = self.weight_quantization_directions(
209 needs_grad_input
210 )
211 # abstract: 量化 GMM-ready [E,K,N] weight
212 weight_quant = self.quantize_weight(
213 gmm_weight,
214 rowwise=weight_rowwise,
215 colwise=weight_colwise,
216 )
217 # abstract: 执行前向 NN GMM，产出真实输出
218 output = self.grouped_matmul(
219 input_quant,
220 weight_quant,
221 layout="NN",
222 group_list=effective_group_list,
223 group_type=0,
224 group_list_type=effective_group_list_type,
225 output_dtype=inputs.dtype,
226 )
228 # The forward view is spent; keep only the direction backward needs.
229 input_quant.update_usage(rowwise=False, colwise=needs_grad_weight)
230 # hook: 只保留 dgrad 需要的 weight 方向——W4A8 覆写
231 self.retain_weight_backward(weight_quant, needs_grad_input)
232 operands = (
233 input_quant if needs_grad_weight else None,
234 weight_quant if needs_grad_input else None,
235 )
236 ctx.quantized_operands_saved = all(
237 operand is None or isinstance(operand, QuantizedTensor)
238 for operand in operands
239 )
240 if ctx.quantized_operands_saved:
241 save_quantized_operands(ctx, *operands, effective_group_list)
242 ctx.input_quant = None
243 ctx.weight_quant = None
244 else:
245 ctx.save_for_backward(effective_group_list)
246 ctx.input_quant, ctx.weight_quant = operands
247 return output
249 def backward(
250 self,
251 ctx: torch.autograd.function.FunctionCtx,
252 grad_output: torch.Tensor,
253 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
254 """Run the complete dgrad/wgrad lifecycle using this strategy's hooks."""
256 needs_grad_input = ctx.needs_input_grad[0]
257 needs_grad_weight = ctx.needs_input_grad[1]
258 if ctx.empty_input:
259 (group_list,) = ctx.saved_tensors
260 grad_input = (
261 torch.zeros(
262 ctx.input_shape,
263 dtype=ctx.input_dtype,
264 device=ctx.input_device,
265 )
266 if needs_grad_input
267 else None
268 )
269 grad_weight = (
270 torch.zeros(
271 ctx.weight_shape,
272 dtype=ctx.weight_dtype,
273 device=ctx.weight_device,
274 )
275 if needs_grad_weight
276 else None
277 )
278 return grad_input, grad_weight
280 if ctx.quantized_operands_saved:
281 group_list, input_quant, weight_quant = restore_quantized_operands(ctx)
282 else:
283 (group_list,) = ctx.saved_tensors
284 input_quant = ctx.input_quant
285 weight_quant = ctx.weight_quant
286 # abstract: 量化 grad_output，保留 backward 需要的两个方向
287 grad_quant = self.quantize_grad_output(
288 grad_output,
289 rowwise=needs_grad_input,
290 colwise=needs_grad_weight,
291 group_list=group_list,
292 group_list_type=ctx.group_list_type,
293 )
294 grad_input = None
295 grad_weight = None
296 if needs_grad_input:
297 # abstract: dgrad——grad_output × 保留的 weight（NT）
298 grad_input = self.grouped_matmul(
299 grad_quant,
300 weight_quant,
301 layout="NT",
302 group_list=group_list,
303 group_type=0,
304 group_list_type=ctx.group_list_type,
305 output_dtype=ctx.input_dtype,
306 )
307 if needs_grad_weight:
308 # abstract: wgrad——保留的 input × grad_output（TN，得 [E,K,O]）
309 grad_weight = self.grouped_matmul(
310 input_quant,
311 grad_quant,
312 layout="TN",
313 group_list=group_list,
314 group_type=2,
315 group_list_type=ctx.group_list_type,
316 output_dtype=ctx.weight_dtype,
317 )
318 # The TN wgrad is GMM-shaped ``[E, K, O]``; transpose back to the
319 # live ``[E, O, K]`` parameter layout.
320 grad_weight = grad_weight.transpose(-2, -1).contiguous()
322 grad_quant.update_usage(rowwise=False, colwise=False)
323 ctx.input_quant = None
324 ctx.weight_quant = None
325 return grad_input, grad_weight
328def validate_grouped_linear_inputs(
329 inputs: torch.Tensor,
330 weight: torch.Tensor,
331 group_list: torch.Tensor,
332 group_list_type: int,
333 *,
334 weight_input_dim: int,
335 name: str,
336) -> None:
337 """Validate shape and grouping invariants shared by native formats."""
339 if inputs.ndim != 2:
340 raise ValueError(
341 f"{name} grouped linear inputs must be two-dimensional, "
342 f"got shape {tuple(inputs.shape)}."
343 )
344 if weight.ndim != 3:
345 raise ValueError(
346 f"{name} grouped linear weight must be three-dimensional, "
347 f"got shape {tuple(weight.shape)}."
348 )
349 if inputs.shape[-1] != weight.shape[weight_input_dim]:
350 raise ValueError(
351 f"{name} grouped linear contracting dimensions differ: "
352 f"inputs={inputs.shape[-1]}, weight={weight.shape[weight_input_dim]}."
353 )
354 if not isinstance(group_list, torch.Tensor) or group_list.ndim != 1:
355 raise ValueError(f"{name} grouped linear group_list must be one-dimensional.")
356 if group_list.shape[0] != weight.shape[0]:
357 raise ValueError(
358 f"{name} grouped linear requires one group per expert: "
359 f"groups={group_list.shape[0]}, experts={weight.shape[0]}."
360 )
361 if group_list_type not in (0, 1):
362 raise ValueError(
363 f"{name} grouped linear group_list_type must be 0 or 1, "
364 f"but got {group_list_type}."
365 )
368class _GroupedLinearFunction(torch.autograd.Function):
369 """Bridge PyTorch autograd callbacks to the strategy selected for the call.
371 Model adapters pass the strategy instance into ``apply`` as an argument;
372 inside the bridge the strategy's own forward/backward lifecycle runs.
373 Because ``Function`` forward/backward are static methods, instance state
374 (quantizer, tile size, ...) can only reach the flow through the ``apply``
375 arguments, so the strategy is carried on ``ctx`` for backward.
376 """
378 @staticmethod
379 def forward(
380 ctx: torch.autograd.function.FunctionCtx,
381 inputs: torch.Tensor,
382 weight: torch.Tensor,
383 group_list: torch.Tensor,
384 strategy: GroupedLinear,
385 group_list_type: int,
386 ) -> torch.Tensor:
387 """Delegate the complete forward lifecycle to the strategy."""
389 return strategy.forward(
390 ctx,
391 inputs,
392 weight,
393 group_list,
394 group_list_type,
395 )
397 @staticmethod
398 def backward(
399 ctx: torch.autograd.function.FunctionCtx,
400 grad_output: torch.Tensor,
401 ) -> tuple[
402 Optional[torch.Tensor],
403 Optional[torch.Tensor],
404 None,
405 None,
406 None,
407 ]:
408 """Delegate the complete backward lifecycle to the strategy."""
410 grad_input, grad_weight = ctx.strategy.backward(ctx, grad_output)
411 return grad_input, grad_weight, None, None, None
414__all__ = [
415 "GroupedLinear",
416 "_GroupedLinearFunction",
417 "validate_grouped_linear_inputs",
418]