Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / base_gmm_func.py: 96%
114 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"""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, their forward weight
20direction, group normalization, and backend calls. The shared template derives
21the dgrad direction and owns weight-view cleanup.
22"""
24from abc import ABC, abstractmethod
25from typing import Optional
27import torch # pylint: disable=forbidden-backend-import
29from hyper_parallel.components.quantization.tensor import QuantizedTensor
30from hyper_parallel.components.quantization.config import (
31 GMM_GROUP_K,
32 GMM_GROUP_M,
33 GROUPED_NDIM,
34 GROUP_LIST_COUNTS,
35 GROUP_LIST_OFFSETS,
36 MATRIX_NDIM,
37)
38from hyper_parallel.components.quantization.functional._saved_quantized import (
39 restore_quantized_operands,
40 save_quantized_operands,
41)
43class GroupedLinear(ABC):
44 """Own the common grouped-linear flow and expose format-specific hooks.
46 The weight contract is fixed to the DeepSeek-V3 live layout
47 ``[E, O, K]`` (contracting ``K`` at ``-1``): the adapter passes the packed
48 source parameters through unchanged. ``forward`` transposes each weight
49 once, inside the flow, to the GMM-ready ``[E, K, N]`` that strategies and
50 operators consume; ``backward`` transposes the wgrad result back to
51 ``[E, O, K]``.
52 """
54 FORMAT_NAME: str
56 def _validate_inputs(
57 self,
58 inputs: torch.Tensor,
59 weight: torch.Tensor,
60 group_list: torch.Tensor,
61 group_list_type: int,
62 ) -> None:
63 """Validate the common live-weight and group metadata contract."""
65 validate_grouped_linear_inputs(
66 inputs,
67 weight,
68 group_list,
69 group_list_type,
70 weight_input_dim=-1,
71 name=self.FORMAT_NAME,
72 )
74 @abstractmethod
75 def normalize_group_list(
76 self,
77 group_list: torch.Tensor,
78 group_list_type: int,
79 ) -> tuple[torch.Tensor, int]:
80 """Adapt grouped metadata to this backend's fixed contract.
82 ``group_list_type=1`` means ``group_list`` contains per-expert token
83 counts; ``group_list_type=0`` means it contains cumulative expert
84 offsets. The returned tensor and type must describe the same
85 representation. MXFP8 and HiFloat8 preserve either representation,
86 while W4A8 strategies convert counts to cumulative offsets and return
87 type 0 because their grouped-matmul path uses the offsets contract.
88 The shared lifecycle passes this pair through forward and saves it for
89 backward, so callers must not assume the returned type is always 0.
90 """
92 @abstractmethod
93 def output_features(self, weight: torch.Tensor) -> int:
94 """Return the output width ``O`` in the live ``[E, O, K]`` layout."""
96 @property
97 @abstractmethod
98 def weight_forward_is_rowwise(self) -> bool:
99 """Declare the weight view required by this strategy's forward NN GMM.
101 Return ``True`` when the forward kernel consumes the row-wise weight
102 view, or ``False`` when it consumes the column-wise view. This is a
103 backend/kernel layout contract, not a user-facing quantization option:
104 the shared lifecycle uses it to build the forward view and, when an
105 input gradient is needed, the opposite view for the dgrad NT GMM. After
106 the forward call, it releases the forward view and retains only the
107 view required by backward.
109 The current strategies use these layouts:
111 * MXFP8 and HiFloat8 use the column-wise view for forward NN GMM.
112 * Native W4A8 uses the row-wise view for forward NN GMM.
113 * Fake W4A8 uses the column-wise view with block size 32, matching its
114 ordinary MXFP8 GMM weight layout.
116 Each strategy must declare its layout explicitly because no single
117 default is valid for all supported kernels; an incorrect view can make
118 the quantized weights incompatible with the selected GMM kernel.
120 Returns:
121 Whether the forward NN GMM consumes the row-wise weight view.
122 """
124 @abstractmethod
125 def quantize_input(
126 self,
127 inputs: torch.Tensor,
128 *,
129 rowwise: bool,
130 colwise: bool,
131 group_list: torch.Tensor,
132 group_list_type: int,
133 ) -> QuantizedTensor:
134 """Create the format-specific activation representation."""
136 @abstractmethod
137 def quantize_weight(
138 self,
139 weight: torch.Tensor,
140 *,
141 rowwise: bool,
142 colwise: bool,
143 ) -> QuantizedTensor:
144 """Create the format-specific weight representation."""
146 @abstractmethod
147 def quantize_grad_output(
148 self,
149 grad_output: torch.Tensor,
150 *,
151 rowwise: bool,
152 colwise: bool,
153 group_list: torch.Tensor,
154 group_list_type: int,
155 ) -> QuantizedTensor:
156 """Create the format-specific gradient-output representation."""
158 @abstractmethod
159 def grouped_matmul(
160 self,
161 left: QuantizedTensor,
162 right: QuantizedTensor,
163 *,
164 layout: str,
165 group_list: torch.Tensor,
166 group_type: int,
167 group_list_type: int,
168 output_dtype: torch.dtype,
169 ) -> torch.Tensor:
170 """Execute one format-specific grouped matrix multiplication."""
172 def forward(
173 self,
174 ctx: torch.autograd.function.FunctionCtx,
175 inputs: torch.Tensor,
176 weight: torch.Tensor,
177 group_list: torch.Tensor,
178 group_list_type: int,
179 ) -> torch.Tensor:
180 """Run the complete forward lifecycle using this strategy's hooks."""
182 # Shared validation gate for live [E, O, K] geometry and group
183 # invariants. Errors use FORMAT_NAME as their prefix.
184 self._validate_inputs(
185 inputs,
186 weight,
187 group_list,
188 group_list_type,
189 )
190 # Prepare the strategy-specific representation once. Keep the tensor
191 # and its type paired: W4A8 converts counts to offsets/type 0, whereas
192 # MXFP8 and HiFloat8 can preserve counts/type 1.
193 effective_group_list, effective_group_list_type = self.normalize_group_list(
194 group_list,
195 group_list_type,
196 )
197 ctx.input_shape = inputs.shape
198 ctx.input_dtype = inputs.dtype
199 ctx.input_device = inputs.device
200 ctx.weight_shape = weight.shape
201 ctx.weight_dtype = weight.dtype
202 ctx.weight_device = weight.device
203 # Backward must interpret the saved group list exactly as forward did.
204 ctx.group_list_type = effective_group_list_type
205 ctx.strategy = self
206 ctx.empty_input = inputs.shape[0] == 0
207 if ctx.empty_input:
208 # Return the output width O even when the input is empty.
209 ctx.save_for_backward(effective_group_list)
210 return inputs.new_empty((0, self.output_features(weight)))
212 # The adapter hands live ``[E, O, K]`` weights; transpose once to the
213 # GMM-ready ``[E, K, N]`` that quantizers and operators consume.
214 gmm_weight = weight.transpose(-2, -1).contiguous()
215 needs_grad_input = inputs.requires_grad
216 needs_grad_weight = weight.requires_grad
217 # Quantize activations: rowwise for forward and colwise for wgrad.
218 input_quant = self.quantize_input(
219 inputs,
220 rowwise=True,
221 colwise=needs_grad_weight,
222 group_list=effective_group_list,
223 group_list_type=effective_group_list_type,
224 )
225 # The forward NN and dgrad NT kernels consume opposite weight views.
226 forward_rowwise = self.weight_forward_is_rowwise
227 # Always build the forward view; build the opposite dgrad view only
228 # when the input gradient needs the weight.
229 if forward_rowwise:
230 weight_rowwise, weight_colwise = True, needs_grad_input
231 else:
232 weight_rowwise, weight_colwise = needs_grad_input, True
233 # Quantize the GMM-ready [E, K, N] weight.
234 weight_quant = self.quantize_weight(
235 gmm_weight,
236 rowwise=weight_rowwise,
237 colwise=weight_colwise,
238 )
239 # Run the forward NN GMM and produce the output.
240 output = self.grouped_matmul(
241 input_quant,
242 weight_quant,
243 layout="NN",
244 group_list=effective_group_list,
245 group_type=GMM_GROUP_M,
246 group_list_type=effective_group_list_type,
247 output_dtype=inputs.dtype,
248 )
249 # Some CANN grouped-matmul kernels promote their result to float32
250 # even when ``output_dtype`` is supplied. The routed expert contract
251 # is the input activation dtype; enforce it here so the next model
252 # operator (for example BF16 RMSNorm) receives a supported pair.
253 if output.dtype != inputs.dtype:
254 output = output.to(dtype=inputs.dtype)
256 # The forward view is spent; keep only the direction backward needs.
257 input_quant.update_usage(rowwise=False, colwise=needs_grad_weight)
258 # Dgrad consumes the opposite view. Release both views when it is unused.
259 if not needs_grad_input:
260 weight_quant.update_usage(rowwise=False, colwise=False)
261 elif forward_rowwise:
262 weight_quant.update_usage(rowwise=False, colwise=True)
263 else:
264 weight_quant.update_usage(rowwise=True, colwise=False)
265 operands = (
266 input_quant if needs_grad_weight else None,
267 weight_quant if needs_grad_input else None,
268 )
269 ctx.quantized_operands_saved = all(
270 operand is None or isinstance(operand, QuantizedTensor)
271 for operand in operands
272 )
273 if ctx.quantized_operands_saved:
274 save_quantized_operands(ctx, *operands, effective_group_list)
275 ctx.input_quant = None
276 ctx.weight_quant = None
277 else:
278 ctx.save_for_backward(effective_group_list)
279 ctx.input_quant, ctx.weight_quant = operands
280 return output
282 def backward(
283 self,
284 ctx: torch.autograd.function.FunctionCtx,
285 grad_output: torch.Tensor,
286 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
287 """Run the complete dgrad/wgrad lifecycle using this strategy's hooks."""
289 needs_grad_input = ctx.needs_input_grad[0]
290 needs_grad_weight = ctx.needs_input_grad[1]
291 if ctx.empty_input:
292 (group_list,) = ctx.saved_tensors
293 grad_input = (
294 torch.zeros(
295 ctx.input_shape,
296 dtype=ctx.input_dtype,
297 device=ctx.input_device,
298 )
299 if needs_grad_input
300 else None
301 )
302 grad_weight = (
303 torch.zeros(
304 ctx.weight_shape,
305 dtype=ctx.weight_dtype,
306 device=ctx.weight_device,
307 )
308 if needs_grad_weight
309 else None
310 )
311 return grad_input, grad_weight
313 if ctx.quantized_operands_saved:
314 group_list, input_quant, weight_quant = restore_quantized_operands(ctx)
315 else:
316 (group_list,) = ctx.saved_tensors
317 input_quant = ctx.input_quant
318 weight_quant = ctx.weight_quant
319 # Quantize grad_output with the views required by backward.
320 grad_quant = self.quantize_grad_output(
321 grad_output,
322 rowwise=needs_grad_input,
323 colwise=needs_grad_weight,
324 group_list=group_list,
325 group_list_type=ctx.group_list_type,
326 )
327 grad_input = None
328 grad_weight = None
329 if needs_grad_input:
330 # dgrad: grad_output × retained weight (NT).
331 grad_input = self.grouped_matmul(
332 grad_quant,
333 weight_quant,
334 layout="NT",
335 group_list=group_list,
336 group_type=GMM_GROUP_M,
337 group_list_type=ctx.group_list_type,
338 output_dtype=ctx.input_dtype,
339 )
340 if needs_grad_weight:
341 # wgrad: retained input × grad_output (TN), producing [E, K, O].
342 grad_weight = self.grouped_matmul(
343 input_quant,
344 grad_quant,
345 layout="TN",
346 group_list=group_list,
347 group_type=GMM_GROUP_K,
348 group_list_type=ctx.group_list_type,
349 output_dtype=ctx.weight_dtype,
350 )
351 # The TN wgrad is GMM-shaped ``[E, K, O]``; transpose back to the
352 # live ``[E, O, K]`` parameter layout.
353 grad_weight = grad_weight.transpose(-2, -1).contiguous()
355 grad_quant.update_usage(rowwise=False, colwise=False)
356 ctx.input_quant = None
357 ctx.weight_quant = None
358 return grad_input, grad_weight
361def validate_grouped_linear_inputs(
362 inputs: torch.Tensor,
363 weight: torch.Tensor,
364 group_list: torch.Tensor,
365 group_list_type: int,
366 *,
367 weight_input_dim: int,
368 name: str,
369) -> None:
370 """Validate shape and grouping invariants shared by native formats."""
372 if inputs.ndim != MATRIX_NDIM:
373 raise ValueError(
374 f"{name} grouped linear inputs must be two-dimensional, "
375 f"got shape {tuple(inputs.shape)}."
376 )
377 if weight.ndim != GROUPED_NDIM:
378 raise ValueError(
379 f"{name} grouped linear weight must be three-dimensional, "
380 f"got shape {tuple(weight.shape)}."
381 )
382 if inputs.shape[-1] != weight.shape[weight_input_dim]:
383 raise ValueError(
384 f"{name} grouped linear contracting dimensions differ: "
385 f"inputs={inputs.shape[-1]}, weight={weight.shape[weight_input_dim]}."
386 )
387 if not isinstance(group_list, torch.Tensor) or group_list.ndim != 1:
388 raise ValueError(f"{name} grouped linear group_list must be one-dimensional.")
389 if group_list.shape[0] != weight.shape[0]:
390 raise ValueError(
391 f"{name} grouped linear requires one group per expert: "
392 f"groups={group_list.shape[0]}, experts={weight.shape[0]}."
393 )
394 if group_list_type not in (GROUP_LIST_OFFSETS, GROUP_LIST_COUNTS):
395 raise ValueError(
396 f"{name} grouped linear group_list_type must be "
397 f"{GROUP_LIST_OFFSETS} or {GROUP_LIST_COUNTS}, "
398 f"but got {group_list_type}."
399 )
402# PyTorch's Function base exposes variadic callbacks; this bridge intentionally
403# specializes them for the grouped-linear ``apply`` contract.
404class _GroupedLinearFunction(torch.autograd.Function): # pylint: disable=abstract-method,arguments-differ
405 """Bridge PyTorch autograd callbacks to the strategy selected for the call.
407 Model adapters pass the strategy instance into ``apply`` as an argument;
408 inside the bridge the strategy's own forward/backward lifecycle runs.
409 Because ``Function`` forward/backward are static methods, instance state
410 (quantizer, tile size, ...) can only reach the flow through the ``apply``
411 arguments, so the strategy is carried on ``ctx`` for backward.
412 """
414 @staticmethod
415 def forward( # pylint: disable=arguments-differ
416 ctx: torch.autograd.function.FunctionCtx,
417 inputs: torch.Tensor,
418 weight: torch.Tensor,
419 group_list: torch.Tensor,
420 strategy: GroupedLinear,
421 group_list_type: int,
422 ) -> torch.Tensor:
423 """Delegate the complete forward lifecycle to the strategy."""
425 return strategy.forward(
426 ctx,
427 inputs,
428 weight,
429 group_list,
430 group_list_type,
431 )
433 @staticmethod
434 def backward( # pylint: disable=arguments-differ
435 ctx: torch.autograd.function.FunctionCtx,
436 grad_output: torch.Tensor,
437 ) -> tuple[
438 Optional[torch.Tensor],
439 Optional[torch.Tensor],
440 None,
441 None,
442 None,
443 ]:
444 """Delegate the complete backward lifecycle to the strategy."""
446 grad_input, grad_weight = ctx.strategy.backward(ctx, grad_output)
447 return grad_input, grad_weight, None, None, None
450__all__ = [
451 "GroupedLinear",
452 "_GroupedLinearFunction",
453 "validate_grouped_linear_inputs",
454]