Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / config.py: 87%
83 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"""Typed low-precision policies and their FQN-independent dtype-scheme catalog."""
17from collections.abc import Mapping
18from dataclasses import dataclass, field
19from typing import Literal, Optional
22# One registry row per accepted dtype: (family, bit width, MX block check).
23# weight_format and act_format must share one family (operators never mix
24# families). The weight/activation bit widths are read directly from these
25# rows to select the native strategy (e.g. mxfp4 weight + mxfp8 act => W4A8).
26# Future dtypes (int/fp/bf, e.g. bf16) are one new row here; hif4 is
27# pre-listed as part of the hifloat family (reserved for its flow).
28_FORMAT_SPECS = {
29 "mxfp4": ("mxfp", 4, True),
30 "mxfp8": ("mxfp", 8, True),
31 "hif4": ("hifloat", 4, False),
32 "hif8": ("hifloat", 8, False),
33}
34_FORMATS = tuple(_FORMAT_SPECS)
37def _format_spec(format_name: str) -> tuple[str, int, bool]:
38 """Look up one dtype format in ``_FORMAT_SPECS``.
40 Drives ``LowPrecisionDtypeScheme.__post_init__`` validation (family /
41 bit-width / MX-block rules). Raises ``ValueError`` on an unregistered
42 format so a bad ``weight_format``/``act_format`` fails at config time.
43 """
45 try:
46 return _FORMAT_SPECS[format_name]
47 except (KeyError, TypeError) as exc:
48 raise ValueError(
49 f"Unsupported low-precision format {format_name!r}; "
50 f"expected {_FORMATS}."
51 ) from exc
54def _format_family(format_name: str) -> str:
55 """Resolve one dtype format to its family name for the factory family gate.
57 Exported for the strategy factory, which must reject cross-family schemes
58 without coupling to the raw ``_FORMAT_SPECS`` tuple layout.
59 """
61 return _format_spec(format_name)[0]
64@dataclass(frozen=True)
65class LowPrecisionDtypeScheme:
66 """One complete low-precision dtype scheme referenced by a plan entry.
67 Defines the quantization dtypes for weights and activations together with
68 the quantization granularity and whether the scheme runs the native
69 operator flow or a future fake-quantize (fake QAT) seam. The native
70 strategy (W4A8/W8A8) is derived downstream from ``weight_format`` /
71 ``act_format`` bit widths, so it is not a separate field here.
72 """
74 # Optional: True = future fake-QAT seam; False (default) = native flow.
75 is_fake_quantize: bool = False
76 weight_format: Literal["mxfp4", "mxfp8", "hif4", "hif8"] = "mxfp8"
77 act_format: Literal["mxfp4", "mxfp8", "hif4", "hif8"] = "mxfp8"
78 block_size: int = 32
80 def __post_init__(self) -> None:
81 """Validate that the scheme names a consistent native/QAT policy.
82 ``weight_format``/``act_format`` must share one format family (the
83 underlying operator never mixes families); their bit widths drive
84 strategy selection downstream. MX formats require a 32/128 block size.
85 """
87 if not isinstance(self.is_fake_quantize, bool):
88 raise ValueError(
89 "LowPrecisionDtypeScheme.is_fake_quantize must be a bool, "
90 f"but got {self.is_fake_quantize!r}."
91 )
92 if isinstance(self.block_size, bool) or not isinstance(self.block_size, int):
93 raise ValueError("LowPrecisionDtypeScheme.block_size must be an integer.")
95 weight_family, weight_bits, weight_is_mx = _format_spec(self.weight_format)
96 act_family, act_bits, act_is_mx = _format_spec(self.act_format)
97 if weight_family != act_family:
98 raise ValueError(
99 "weight_format and act_format must belong to the same format "
100 "family (the underlying operator never mixes families); got "
101 f"weight={self.weight_format!r}/act={self.act_format!r}."
102 )
103 if (weight_is_mx or act_is_mx) and self.block_size not in (32, 128):
104 raise ValueError(
105 "MX low-precision block_size must be 32 or 128, "
106 f"but got {self.block_size}."
107 )
108 if (
109 not self.is_fake_quantize
110 and weight_family == "mxfp"
111 and weight_bits == 8
112 and self.block_size != 32
113 ):
114 raise ValueError(
115 "native w8a8 mxfp8 scheme requires block_size=32 because the "
116 "current MXFP8 operator contract uses 32-element blocks."
117 )
120@dataclass
121class LowPrecisionConfig:
122 """Global switch plus a catalog of named dtype schemes.
124 A dtype scheme is selected by the surrounding ``PlanOverride`` entry.
125 Matching module FQNs and selecting a dtype scheme are intentionally
126 separate concerns: this object is only the low-precision catalog.
128 ``format``/``scaling`` retain the pre-catalog configuration surface. They
129 are converted to one implicit scheme when no named catalog is supplied,
130 so existing callers remain valid while mixed W8A8/W4A8 plans use
131 ``dtype_schemes`` explicitly.
132 """
134 enabled: bool = False
135 format: Literal["mxfp8_e4m3", "mxfp4_e2m1", "hif8"] = "mxfp8_e4m3"
136 scaling: Literal["mx_block", "current"] = "mx_block"
137 dtype_schemes: dict[str, LowPrecisionDtypeScheme] = field(default_factory=dict)
138 default_dtype_scheme: Optional[str] = None
140 def __post_init__(self) -> None:
141 """Validate catalog keys and the optional default dtype scheme."""
143 if not isinstance(self.enabled, bool):
144 raise ValueError("LowPrecisionConfig.enabled must be a bool.")
145 legacy_pairs = {
146 ("mxfp8_e4m3", "mx_block"),
147 ("mxfp4_e2m1", "mx_block"),
148 ("hif8", "current"),
149 }
150 if (self.format, self.scaling) not in legacy_pairs:
151 raise ValueError(
152 "Unsupported low-precision format/scaling combination "
153 f"{self.format!r}/{self.scaling!r}."
154 )
156 if not isinstance(self.dtype_schemes, Mapping):
157 raise ValueError(
158 "LowPrecisionConfig.dtype_schemes must be a mapping of "
159 "non-empty names to dtype schemes."
160 )
161 typed_schemes = {}
162 for name, scheme in self.dtype_schemes.items():
163 if isinstance(scheme, Mapping):
164 try:
165 scheme = LowPrecisionDtypeScheme(**dict(scheme))
166 except (TypeError, ValueError) as exc:
167 raise ValueError(
168 f"Invalid low-precision dtype scheme {name!r}: {exc}"
169 ) from exc
170 typed_schemes[name] = scheme
171 self.dtype_schemes = typed_schemes
173 if any(
174 not isinstance(name, str) or not name
175 or not isinstance(scheme, LowPrecisionDtypeScheme)
176 for name, scheme in self.dtype_schemes.items()
177 ):
178 raise ValueError(
179 "LowPrecisionConfig.dtype_schemes must map non-empty names to "
180 "LowPrecisionDtypeScheme objects."
181 )
182 if self.default_dtype_scheme is not None:
183 if not isinstance(self.default_dtype_scheme, str) or not self.default_dtype_scheme:
184 raise ValueError(
185 "LowPrecisionConfig.default_dtype_scheme must be a dtype scheme name."
186 )
187 if self.default_dtype_scheme not in self.dtype_schemes:
188 raise ValueError(
189 f"LowPrecisionConfig.default_dtype_scheme "
190 f"{self.default_dtype_scheme!r} is not present in dtype_schemes."
191 )
193 def resolve_dtype_scheme(
194 self, dtype_scheme_name: Optional[str] = None,
195 ) -> LowPrecisionDtypeScheme:
196 """Map a dtype-scheme *name* to the scheme object used downstream.
198 The input is either ``plan_overrides[].low_precision_dtype_scheme`` or
199 ``default_dtype_scheme``; the returned ``LowPrecisionDtypeScheme`` is
200 what the strategy factory and adapters actually consume. Only resolves
201 catalog names -- which modules are matched is decided by the plan.
202 """
204 selected = dtype_scheme_name or self.default_dtype_scheme
205 if selected is not None:
206 if not isinstance(selected, str) or not selected:
207 raise ValueError(
208 "low-precision dtype scheme name must be a non-empty string."
209 )
210 try:
211 return self.dtype_schemes[selected]
212 except KeyError as exc:
213 raise ValueError(
214 f"Unknown low-precision dtype scheme {selected!r}; "
215 f"available dtype schemes: {sorted(self.dtype_schemes)}."
216 ) from exc
217 if self.dtype_schemes:
218 # A catalog exists but nothing was selected: fail loudly instead of
219 # silently running the model unquantized.
220 raise ValueError(
221 "LowPrecisionConfig.dtype_schemes is configured, but the replacement "
222 "plan did not select low_precision_dtype_scheme and no "
223 "default_dtype_scheme was declared."
224 )
225 legacy_scheme = {
226 ("mxfp8_e4m3", "mx_block"): LowPrecisionDtypeScheme(),
227 ("mxfp4_e2m1", "mx_block"): LowPrecisionDtypeScheme(
228 weight_format="mxfp4",
229 act_format="mxfp8",
230 ),
231 ("hif8", "current"): LowPrecisionDtypeScheme(
232 weight_format="hif8",
233 act_format="hif8",
234 ),
235 }
236 return legacy_scheme[(self.format, self.scaling)]
238 def to_dict(self) -> dict[str, object]:
239 """Serialize legacy fields and the optional named scheme catalog."""
241 result: dict[str, object] = {
242 "enabled": self.enabled,
243 "format": self.format,
244 "scaling": self.scaling,
245 }
246 if self.dtype_schemes:
247 result["dtype_schemes"] = {
248 name: {
249 "is_fake_quantize": scheme.is_fake_quantize,
250 "weight_format": scheme.weight_format,
251 "act_format": scheme.act_format,
252 "block_size": scheme.block_size,
253 }
254 for name, scheme in self.dtype_schemes.items()
255 }
256 if self.default_dtype_scheme is not None:
257 result["default_dtype_scheme"] = self.default_dtype_scheme
258 return result
261__all__ = [
262 "LowPrecisionConfig",
263 "LowPrecisionDtypeScheme",
264]