Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / quantizers / hifloat8.py: 42%
55 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"""Pangu-compatible HiFloat8 current-scaling quantizer."""
17from typing import Any, Optional
19import torch # pylint: disable=forbidden-backend-import
21from hyper_parallel.components.quantization.config import (
22 GROUPED_NDIM,
23 GROUP_LIST_COUNTS,
24 GROUP_LIST_OFFSETS,
25 MATRIX_NDIM,
26)
27from hyper_parallel.components.quantization.ops.npu_hifloat8 import (
28 HiFloat8NpuOps,
29 _get_hifloat8_npu_ops,
30)
31from hyper_parallel.components.quantization.quantizers.base import Quantizer
32from hyper_parallel.components.quantization.tensor.hifloat8_tensor import (
33 HiFloat8Tensor,
34)
35from hyper_parallel.components.quantization.tensor.quantized_tensor import (
36 QuantizedTensor,
37)
39INPUT_WEIGHT_FORMAT_MAX = 15.0
40GRADIENT_FORMAT_MAX = 224.0
41_SUPPORTED_FORMAT_MAX = (INPUT_WEIGHT_FORMAT_MAX, GRADIENT_FORMAT_MAX)
44class HiFloat8Quantizer(Quantizer):
45 """Quantize one HiFloat8 role with stateless current tensorwise scaling."""
47 def __init__(
48 self,
49 *,
50 fp8_max: float,
51 npu_ops: Optional[HiFloat8NpuOps] = None,
52 ) -> None:
53 """Create one frozen-role quantizer compatible with Pangu's structure."""
55 if float(fp8_max) not in _SUPPORTED_FORMAT_MAX:
56 raise ValueError(
57 "HiFloat8 current scaling supports format max 15 or 224, "
58 f"got {fp8_max}."
59 )
60 self.fp8_max = float(fp8_max)
61 self._npu_ops = npu_ops
63 @property
64 def npu_ops(self) -> HiFloat8NpuOps:
65 """Return the injected or lazily resolved NPU operator adapter."""
67 if self._npu_ops is None:
68 self._npu_ops = _get_hifloat8_npu_ops()
69 return self._npu_ops
71 @property
72 def quant_dtype(self) -> Any:
73 """Return the runtime-frozen logical HiFloat8 dtype identifier."""
75 return self.npu_ops.quant_dtype
77 @staticmethod
78 def _validate_input(
79 tensor: torch.Tensor,
80 *,
81 group_list: Optional[torch.Tensor],
82 group_list_type: int,
83 rowwise: bool,
84 colwise: bool,
85 ) -> None:
86 """Validate tensor dtype, rank, directions, and group metadata."""
87 if not isinstance(tensor, torch.Tensor) or isinstance(
88 tensor,
89 QuantizedTensor,
90 ):
91 raise TypeError(
92 "HiFloat8Quantizer expects a high-precision torch.Tensor."
93 )
94 if tensor.dtype not in (torch.float16, torch.bfloat16):
95 raise TypeError(
96 "HiFloat8 current quantization requires float16 or bfloat16 "
97 f"input, got {tensor.dtype}."
98 )
99 if tensor.ndim not in (MATRIX_NDIM, GROUPED_NDIM):
100 raise ValueError(
101 "HiFloat8 quantization requires a matrix or grouped weight, "
102 "got shape "
103 f"{tuple(tensor.shape)}."
104 )
105 if not rowwise and not colwise:
106 raise ValueError(
107 "HiFloat8Quantizer requires rowwise=True or colwise=True."
108 )
109 if group_list_type not in (GROUP_LIST_OFFSETS, GROUP_LIST_COUNTS):
110 raise ValueError(
111 "HiFloat8Quantizer group_list_type must be 0 or 1, "
112 f"got {group_list_type}."
113 )
114 if group_list is not None:
115 if tensor.ndim != MATRIX_NDIM:
116 raise ValueError(
117 "Grouped HiFloat8 quantization requires a 2D operand."
118 )
119 if not isinstance(group_list, torch.Tensor) or group_list.ndim != 1:
120 raise ValueError(
121 "HiFloat8Quantizer group_list must be one-dimensional."
122 )
123 if group_list.dtype != torch.int64:
124 raise TypeError(
125 "HiFloat8Quantizer group_list must use torch.int64."
126 )
128 def quantize(
129 self,
130 tensor: torch.Tensor,
131 *,
132 group_list: Optional[torch.Tensor] = None,
133 group_list_type: int = GROUP_LIST_OFFSETS,
134 rowwise: bool = False,
135 colwise: bool = False,
136 ) -> HiFloat8Tensor:
137 """Quantize Dense or grouped operands with current scaling.
139 Two-dimensional activation/gradient operands use one tensorwise scale.
140 GMM callers expand that scalar to one identical entry per expert.
141 Three-dimensional expert weights use one scale per leading expert.
142 """
144 self._validate_input(
145 tensor,
146 group_list=group_list,
147 group_list_type=group_list_type,
148 rowwise=rowwise,
149 colwise=colwise,
150 )
151 if tensor.ndim == GROUPED_NDIM:
152 flattened = tensor.reshape(tensor.shape[0], -1)
153 data, scale = self.npu_ops.dynamic_quant(
154 flattened,
155 dst_type_max=self.fp8_max,
156 quant_mode="pertoken",
157 )
158 data = data.reshape_as(tensor)
159 else:
160 data, scale = self.npu_ops.dynamic_quant(
161 tensor,
162 dst_type_max=self.fp8_max,
163 quant_mode="pertensor",
164 )
165 scale = torch.where(
166 torch.isfinite(scale) & (scale > 0),
167 scale,
168 torch.ones_like(scale),
169 )
170 if group_list is not None and scale.numel() == 1:
171 scale = scale.reshape(1).expand(group_list.numel()).contiguous()
172 return HiFloat8Tensor(
173 shape=tensor.shape,
174 dtype=tensor.dtype,
175 quantizer=self,
176 row_data=data if rowwise else None,
177 row_scale=scale if rowwise else None,
178 col_data=data if colwise else None,
179 col_scale=scale if colwise else None,
180 device=tensor.device,
181 )
184__all__ = [
185 "GRADIENT_FORMAT_MAX",
186 "HiFloat8Quantizer",
187 "INPUT_WEIGHT_FORMAT_MAX",
188]