codekingpro/portable-devtools
115k
1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for
4# license information.
5# --------------------------------------------------------------------------
6import logging
7from typing import Any
8
9import numpy as np
10import onnx
11import onnx.numpy_helper
12
13try:
14 from onnx.reference.op_run import to_array_extended
15except ImportError:
16 # old version of onnx.
17 to_array_extended = None
18
19from .calibrate import TensorData
20from .onnx_model import ONNXModel
21from .quant_utils import (
22 DEQUANT_OP_NAME,
23 ONNX_TYPE_TO_NP_TYPE,
24 QUANT_OP_NAME,
25 TENSOR_NAME_QUANT_SUFFIX,
26 find_by_name,
27 get_opset_version,
28 model_has_infer_metadata,
29 normalize_axis,
30 pack_bytes_to_4bit,
31 quantize_data,
32 quantize_nparray,
33 save_and_reload_model_with_shape_infer,
34 tensor_proto_to_array,
35)
36from .tensor_quant_overrides import TensorQuantOverridesHelper
37
38
39class QuantizationParams:
40 def __init__(self, **data: dict[str, Any]):
41 self.data = {}
42 for k, v in data.items():
43 if not isinstance(k, str):
44 raise TypeError(f"Keys must be strings not {type(k)} for k={k!r}.")
45 if k != "axis" and not isinstance(v, (int, str, np.ndarray, float)):
46 raise TypeError(f"Values must be numpy arrays, int, float, str not {type(v)} for k={k!r}.")
47 if k == "axis" and not isinstance(v, int) and v is not None:
48 raise TypeError(f"Axis value must be an int or None, not {type(v)}.")
49 if k == "scale" and v.dtype not in (np.float32, np.float16):
50 raise ValueError(f"scale must a float32 or float16 numpy element but is {v.dtype} for k={k!r}")
51 self.data[k] = v
52
53 def get(self, key, default_value=None):
54 return self.data.get(key, default_value)
55
56 def __iter__(self):
57 yield from self.data
58
59 def __getitem__(self, key):
60 return self.data[key]
61
62 def __setitem__(self, key, value):
63 self.data[key] = value
64
65 def __len__(self):
66 return len(self.data)
67
68
69class BaseQuantizer:
70 def __init__(
71 self,
72 model,
73 per_channel,
74 reduce_range,
75 weight_qType,
76 activation_qType,
77 tensors_range,
78 nodes_to_quantize,
79 nodes_to_exclude,
80 op_types_to_quantize,
81 extra_options=None,
82 ):
83 if not model_has_infer_metadata(model):
84 model = save_and_reload_model_with_shape_infer(model)
85 self.value_infos = {vi.name: vi for vi in model.graph.value_info}
86 self.value_infos.update({ot.name: ot for ot in model.graph.output})
87 self.value_infos.update({it.name: it for it in model.graph.input})
88
89 self.model = ONNXModel(model)
90 self.opset_version = get_opset_version(model)
91 self.per_channel = per_channel # weight-pack per channel
92 self.reduce_range = reduce_range
93
94 self.extra_options = extra_options if extra_options else {}
95 self.enable_subgraph_quantization = (
96 "EnableSubgraph" in self.extra_options and self.extra_options["EnableSubgraph"]
97 )
98 self.parent = None
99 self.force_quantize_no_input_check = (
100 "ForceQuantizeNoInputCheck" in self.extra_options and self.extra_options["ForceQuantizeNoInputCheck"]
101 )
102
103 # If user does not explicitly set "WeightSymmetric", then the weight's quantization type determines
104 # the symmetry (i.e., signed integer types will use symmetric quantization). See `def is_weight_symmetric()`
105 self._is_weight_symmetric: bool | None = self.extra_options.get("WeightSymmetric", None)
106 self.is_activation_symmetric = self.extra_options.get("ActivationSymmetric", False)
107 self.min_real_range = self.extra_options.get("MinimumRealRange")
108
109 self.activation_qType = getattr(activation_qType, "tensor_type", activation_qType)
110 self.weight_qType = getattr(weight_qType, "tensor_type", weight_qType)
111
112 """
113 Dictionary specifying the min and max values for tensors. It has following format:
114 {
115 "param_name": [min, max]
116 }
117 example:
118 {
119 'Conv_3:0': [np.float32(0), np.float32(0.5)],
120 'Conv_4:0': [np.float32(1), np.float32(3.5)]
121 }
122 """
123 if tensors_range is not None and any(not isinstance(t, TensorData) for t in tensors_range.values()):
124 raise TypeError(
125 f"tensors_range contains unexpected types { {type(v) for v in tensors_range.values()} }, not TensorData."
126 )
127 self.tensors_range = tensors_range
128 self.nodes_to_quantize = nodes_to_quantize # specific nodes to quantize
129 self.nodes_to_exclude = nodes_to_exclude # specific nodes to exclude
130 self.op_types_to_quantize = op_types_to_quantize
131
132 # Get tensor-level quantization overrides and ensure they are valid.
133 self.tensor_quant_overrides = TensorQuantOverridesHelper(self.extra_options.get("TensorQuantOverrides", {}))
134
135 self.initializers = {initzer.name: initzer for initzer in self.model.initializer()}
136 overrides_valid, overrides_err = self.tensor_quant_overrides.is_valid(
137 self.initializers, self.value_infos.keys(), activation_qType
138 )
139 if not overrides_valid:
140 raise ValueError(overrides_err)
141
142 self.tensor_quant_override_qtypes = self.tensor_quant_overrides.get_quant_types()
143
144 def is_weight_symmetric(self, weight_quant_type: onnx.TensorProto.DataType) -> bool:
145 if self._is_weight_symmetric is not None:
146 return self._is_weight_symmetric # Return value explicitly set by user.
147 return weight_quant_type in (
148 onnx.TensorProto.INT4,
149 onnx.TensorProto.INT8,
150 onnx.TensorProto.INT16,
151 onnx.TensorProto.FLOAT8E4M3FN,
152 )
153
154 def quantize_model(self):
155 raise NotImplementedError
156
157 def is_input_a_initializer(self, input_name):
158 initializer = find_by_name(input_name, self.model.initializer())
159 return initializer is not None
160
161 def is_per_channel(self):
162 return self.per_channel
163
164 def is_valid_quantize_weight(self, weight_name):
165 weight = find_by_name(weight_name, self.model.initializer())
166 if weight is not None:
167 return weight.data_type in (onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16)
168 if (not self.enable_subgraph_quantization) or (self.parent is None):
169 return False
170 return self.parent.is_valid_quantize_weight(weight_name)
171
172 def should_quantize_node(self, node):
173 if (
174 self.nodes_to_quantize is not None
175 and len(self.nodes_to_quantize) != 0
176 and node.name not in self.nodes_to_quantize
177 ):
178 return False
179
180 if node.op_type not in self.op_types_to_quantize:
181 return False
182
183 if node.op_type in (DEQUANT_OP_NAME, QUANT_OP_NAME):
184 return False
185
186 if self.nodes_to_exclude is not None and node.name in self.nodes_to_exclude:
187 return False
188
189 return True
190
191 def quantize_bias_static_impl(self, bias_name, input_scale, weight_scale, beta=1.0):
192 """
193 Quantized the bias. Zero Point == 0 and Scale == Input_Scale * Weight_Scale
194 """
195
196 # get bias
197 bias_initializer = find_by_name(bias_name, self.model.initializer())
198 bias_data = tensor_proto_to_array(bias_initializer)
199 quantized_bias_name = bias_name + TENSOR_NAME_QUANT_SUFFIX
200
201 # quantize bias
202 if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
203 data = np.asarray(bias_data)
204 if data.dtype == np.float16:
205 node_qtype = onnx.TensorProto.FLOAT16
206 elif data.dtype == np.float32:
207 node_qtype = onnx.TensorProto.FLOAT
208 else:
209 raise TypeError(f"Only float16 or float32 are supported with float 8 but bias dtype is {data.dtype}.")
210 quantized_data = data.astype(np.float32)
211 bias_scale = np.array([1], dtype=quantized_data.dtype)
212 bias_scale_data = bias_scale.reshape(-1)
213 packed_bias_initializer = onnx.numpy_helper.from_array(quantized_data, quantized_bias_name)
214 self.model.initializer_extend([packed_bias_initializer])
215 node_type = "Cast"
216 else:
217 # calculate scale for bias
218 # TODO: This formula should be explained including why the scale is not estimated for the bias as well.
219 bias_scale = input_scale * weight_scale * beta
220
221 # Quantize by dividing by bias_scale
222 quantized_data = np.asarray(bias_data, dtype=np.float64) / np.asarray(bias_scale, dtype=np.float64)
223 quantized_data = quantized_data.round()
224
225 # Clip quantized data to the range of a int32
226 int32_min = np.float64(np.iinfo(np.int32).min)
227 int32_max = np.float64(np.iinfo(np.int32).max)
228 if np.any(quantized_data < int32_min) or np.any(quantized_data > int32_max):
229 logging.warning(
230 f"Quantized bias `{bias_name}` exceeds the range of a int32. The bias scale is too small."
231 )
232
233 quantized_data = np.clip(quantized_data, int32_min, int32_max).astype(np.int32)
234
235 # update bias initializer
236 bias_np_data = np.asarray(quantized_data, dtype=np.int32).reshape(bias_initializer.dims)
237 packed_bias_initializer = onnx.numpy_helper.from_array(bias_np_data, quantized_bias_name)
238 self.model.initializer_extend([packed_bias_initializer])
239
240 # Bias's scale dtype should match the original bias data's unquantized type (float32 or float16).
241 bias_scale_data = np.asarray(bias_scale, dtype=bias_data.dtype).reshape(-1)
242 node_type = "DequantizeLinear"
243 node_qtype = self.weight_qType
244
245 # update scale initializer
246 quantized_bias_scale_name = quantized_bias_name + "_scale"
247 packed_bias_scale_initializer = onnx.numpy_helper.from_array(bias_scale_data, quantized_bias_scale_name)
248 self.model.initializer_extend([packed_bias_scale_initializer])
249
250 # update zero initializer
251 if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
252 tensor_type = self.weight_qType
253 else:
254 tensor_type = onnx.TensorProto.INT32
255
256 quantized_bias_zp_name = quantized_bias_name + "_zero_point"
257 if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
258 packed_bias_zp_initializer = onnx.helper.make_tensor(quantized_bias_zp_name, self.weight_qType, [1], [0.0])
259 elif bias_scale.size > 1:
260 bias_zp_data = np.zeros(bias_scale.shape, dtype=np.int32).reshape(-1)
261 packed_bias_zp_initializer = onnx.numpy_helper.from_array(bias_zp_data, quantized_bias_zp_name)
262 else:
263 packed_bias_zp_initializer = onnx.helper.make_tensor(quantized_bias_zp_name, tensor_type, [], [0])
264 self.model.initializer_extend([packed_bias_zp_initializer])
265
266 return (
267 quantized_bias_name,
268 quantized_bias_scale_name,
269 quantized_bias_zp_name,
270 bias_scale_data,
271 node_type,
272 node_qtype,
273 )
274
275 def quantize_initializer_impl(self, weight, qType, reduce_range=False, keep_float_weight=False):
276 """
277 :param weight: TensorProto initializer
278 :param qType: type to quantize to
279 :param keep_float_weight: Whether to quantize the weight. In some cases, we only want to qunatize scale and zero point.
280 If keep_float_weight is False, quantize the weight, or don't quantize the weight.
281 :return: quantized weight name, zero point name, scale name
282 """
283 # TODO(adrianlizarraga): This function is now only used by onnx_quantizer.py, so move it there.
284 q_weight_name = weight.name + TENSOR_NAME_QUANT_SUFFIX
285 zp_name = weight.name + "_zero_point"
286 scale_name = weight.name + "_scale"
287
288 # Quantize weight data. Use quantization overrides if provided by the user.
289 weight_data = tensor_proto_to_array(weight)
290 quant_overrides = self.tensor_quant_overrides.get_per_tensor_overrides(weight.name, default_val={})
291 if "quant_type" in quant_overrides:
292 qType = quant_overrides["quant_type"].tensor_type # noqa: N806
293
294 if "scale" in quant_overrides and "zero_point" in quant_overrides:
295 zero_point = np.array(quant_overrides["zero_point"], dtype=ONNX_TYPE_TO_NP_TYPE[qType])
296 scale = np.array(quant_overrides["scale"])
297 q_weight_data = quantize_nparray(qType, weight_data.flatten(), scale, zero_point)
298 assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
299 assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
300 f"Unexpected dtype {zero_point.dtype}"
301 )
302 assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
303
304 else:
305 symmetric = self.is_weight_symmetric(qType) if qType == self.weight_qType else self.is_activation_symmetric
306 zero_point, scale, q_weight_data = quantize_data(
307 weight_data.flatten(),
308 qType,
309 quant_overrides.get("symmetric", symmetric),
310 reduce_range=quant_overrides.get("reduce_range", self.reduce_range and reduce_range),
311 min_real_range=self.min_real_range,
312 rmin_override=quant_overrides.get("rmin"),
313 rmax_override=quant_overrides.get("rmax"),
314 )
315
316 assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
317 assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
318 f"Unexpected dtype {zero_point.dtype}"
319 )
320 assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
321
322 scale_dtype = weight.data_type
323 scale_initializer = onnx.helper.make_tensor(scale_name, scale_dtype, [], scale.reshape((-1,)).tolist())
324 zero_initializer = onnx.helper.make_tensor(zp_name, qType, [], zero_point.reshape((-1,)).tolist())
325 self.model.initializer_extend([scale_initializer, zero_initializer])
326
327 if not keep_float_weight:
328 if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
329 q_weight_initializer = onnx.TensorProto()
330 q_weight_initializer.data_type = self.weight_qType
331 q_weight_initializer.dims.extend(weight.dims)
332 q_weight_initializer.name = q_weight_name
333 # Do not remove .flatten().copy() numpy is not clear about data persistence.
334 q_weight_initializer.raw_data = q_weight_data.flatten().copy().tobytes()
335 if to_array_extended is not None:
336 # This test should not be needed but it helped catch some issues
337 # with data persistence and tobytes.
338 check = to_array_extended(q_weight_initializer)
339 if check.shape != weight_data.shape or check.tobytes() != q_weight_data.tobytes():
340 raise RuntimeError(
341 f"The initializer of shape {weight_data.shape} could not be created, expecting "
342 f"{q_weight_data.tobytes()[:10]}, got {check.tobytes()[:10]} and shape={weight.shape}"
343 f"\nraw={str(q_weight_initializer)[:200]}."
344 )
345 elif qType in (onnx.TensorProto.INT4, onnx.TensorProto.UINT4):
346 if q_weight_data.dtype not in (np.int8, np.uint8):
347 raise RuntimeError(
348 f"Quantized weights for {q_weight_name} must be 8-bit before packing as 4-bit values."
349 )
350
351 # We do not use onnx.helper.pack_float32_to_4bit() due to performance.
352 # This can be the difference between a large model taking 30 minutes to quantize vs 5 minutes.
353 packed_data = bytes(pack_bytes_to_4bit(q_weight_data.tobytes()))
354
355 # We only use onnx.helper.make_tensor with raw data due to bug: https://github.com/onnx/onnx/pull/6161
356 q_weight_initializer = onnx.helper.make_tensor(q_weight_name, qType, weight.dims, packed_data, raw=True)
357 else:
358 q_weight_data = np.asarray(q_weight_data, dtype=onnx.helper.tensor_dtype_to_np_dtype(qType)).reshape(
359 weight.dims
360 )
361 q_weight_initializer = onnx.numpy_helper.from_array(q_weight_data, q_weight_name)
362 self.model.initializer_extend([q_weight_initializer])
363
364 return q_weight_name, zp_name, scale_name
365
366 def quantize_weight_per_channel_impl(
367 self,
368 weight_name,
369 weight_qType,
370 channel_axis,
371 reduce_range=True,
372 keep_float_weight=False,
373 ):
374 # TODO(adrianlizarraga): This function is now only used by onnx_quantizer.py, so move it there.
375 initializer = find_by_name(weight_name, self.model.initializer())
376 if initializer is None:
377 raise ValueError("{} is not an initializer", weight_name)
378
379 weights = tensor_proto_to_array(initializer)
380 weights_rank = len(weights.shape)
381 is_axis_valid, axis_norm = normalize_axis(channel_axis, weights_rank)
382 if not is_axis_valid:
383 raise ValueError(
384 f"Weight {weight_name} has a per-channel axis with value {channel_axis} that is "
385 f"out-of-bounds for rank {weights_rank}"
386 )
387
388 channel_axis = axis_norm
389 channel_count = weights.shape[channel_axis]
390 quant_overrides_for_channels = self.tensor_quant_overrides.get_per_channel_overrides(
391 weight_name, default_val=[{"axis": channel_axis}]
392 )
393
394 num_channel_overrides = len(quant_overrides_for_channels)
395 if num_channel_overrides != 1 and num_channel_overrides != channel_count:
396 raise ValueError(
397 f"Per-channel tensor quantization overrides for {weight_name} must have "
398 f"either 1 or {channel_count} elements in the list of dictionaries."
399 )
400
401 is_axis_override_valid, axis_override = normalize_axis(quant_overrides_for_channels[0]["axis"], weights_rank)
402 if not is_axis_override_valid or axis_override != channel_axis:
403 raise ValueError(
404 f"Tensor quantization overrides for {weight_name} specify an unexpected axis. "
405 f"Expected {channel_axis}, but got {quant_overrides_for_channels[0]['axis']}."
406 )
407
408 # If user provides per-channel quantization overrides, all channels must use the same quant_type,
409 # axis, symmetric, and reduce_range values. So, just use the first channel's values.
410 if "quant_type" in quant_overrides_for_channels[0]:
411 weight_qType = quant_overrides_for_channels[0]["quant_type"].tensor_type # noqa: N806
412
413 symmetric = quant_overrides_for_channels[0].get("symmetric", self.is_weight_symmetric(weight_qType))
414 reduce_range = quant_overrides_for_channels[0].get("reduce_range", self.reduce_range and reduce_range)
415 zero_point_list = []
416 scale_list = []
417 quantized_per_channel_data_list = []
418 weights_shape = list(weights.shape)
419 reshape_dims = list(weights_shape) # deep copy
420 reshape_dims[channel_axis] = 1 # only one per channel for reshape
421 for i in range(channel_count):
422 per_channel_data = weights.take(i, channel_axis)
423 channel_override_index = i if i < num_channel_overrides else 0
424 channel_quant_overrides = quant_overrides_for_channels[channel_override_index]
425
426 if "scale" in channel_quant_overrides and "zero_point" in channel_quant_overrides:
427 zero_point = np.array(channel_quant_overrides["zero_point"], dtype=ONNX_TYPE_TO_NP_TYPE[weight_qType])
428 scale = np.array(channel_quant_overrides["scale"])
429 quantized_per_channel_data = quantize_nparray(
430 weight_qType, per_channel_data.flatten(), scale, zero_point
431 )
432 assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
433 assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
434 f"Unexpected dtype {zero_point.dtype}"
435 )
436 assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
437 assert isinstance(quantized_per_channel_data, np.ndarray), (
438 f"Unexpected type {type(quantized_per_channel_data)}"
439 )
440
441 else:
442 zero_point, scale, quantized_per_channel_data = quantize_data(
443 per_channel_data.flatten(),
444 weight_qType,
445 symmetric,
446 reduce_range=reduce_range,
447 min_real_range=self.min_real_range,
448 rmin_override=channel_quant_overrides.get("rmin"),
449 rmax_override=channel_quant_overrides.get("rmax"),
450 )
451
452 assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
453 assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
454 f"Unexpected dtype {zero_point.dtype}"
455 )
456 assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
457 assert isinstance(quantized_per_channel_data, np.ndarray), (
458 f"Unexpected type {type(quantized_per_channel_data)}"
459 )
460
461 zero_point_list.append(zero_point)
462 scale_list.append(scale)
463 quantized_per_channel_data_list.append(np.asarray(quantized_per_channel_data).reshape(reshape_dims))
464
465 # combine per_channel_data into one
466 quantized_weights = np.concatenate(quantized_per_channel_data_list, channel_axis)
467 q_weight_name = weight_name + TENSOR_NAME_QUANT_SUFFIX
468 zp_name = weight_name + "_zero_point"
469 scale_name = weight_name + "_scale"
470
471 # Update packed weight, zero point, and scale initializers
472 zero_scale_shape = [initializer.dims[channel_axis]]
473 scale_initializer = onnx.helper.make_tensor(
474 scale_name, initializer.data_type, zero_scale_shape, np.hstack(scale_list).tolist()
475 )
476 zero_initializer = onnx.helper.make_tensor(
477 zp_name, weight_qType, zero_scale_shape, np.hstack(zero_point_list).tolist()
478 )
479
480 self.model.initializer_extend([scale_initializer, zero_initializer])
481
482 if not keep_float_weight:
483 if weight_qType in (onnx.TensorProto.INT4, onnx.TensorProto.UINT4):
484 if quantized_weights.dtype not in (np.int8, np.uint8):
485 raise RuntimeError(
486 f"Quantized weights for {q_weight_name} must be 8-bit before packing as 4-bit values."
487 )
488
489 # We do not use onnx.helper.pack_float32_to_4bit() due to performance.
490 # This can be the difference between a large model taking 30 minutes to quantize vs 5 minutes.
491 packed_data = bytes(pack_bytes_to_4bit(quantized_weights.tobytes()))
492
493 # We only use onnx.helper.make_tensor with raw data due to bug: https://github.com/onnx/onnx/pull/6161
494 q_weight_initializer = onnx.helper.make_tensor(
495 q_weight_name, weight_qType, weights_shape, packed_data, raw=True
496 )
497 self.model.initializer_extend([q_weight_initializer])
498 else:
499 quantized_weights = np.asarray(
500 quantized_weights,
501 dtype=onnx.helper.tensor_dtype_to_np_dtype(weight_qType),
502 ).reshape(initializer.dims)
503 q_weight_initializer = onnx.numpy_helper.from_array(quantized_weights, q_weight_name)
504 self.model.initializer_extend([q_weight_initializer])
505
506 return q_weight_name, zp_name, scale_name
507
508 def adjust_tensor_ranges(self):
509 if self.tensors_range is None:
510 return
511
512 for node in self.model.nodes():
513 # adjust tensor_ranges for input of Clip and Relu node
514 if node.op_type in ["Clip", "Relu"]:
515 if not self.should_quantize_node(node):
516 continue
517 if len(self.model.input_name_to_nodes()[node.input[0]]) != 1:
518 continue
519 if node.input[0] not in self.tensors_range or node.output[0] not in self.tensors_range:
520 continue
521 td = self.tensors_range[node.output[0]]
522 if not isinstance(td, TensorData):
523 raise TypeError(f"Unexpected type {type(td)} for {node.output[0]!r}.")
524 self.tensors_range[node.input[0]] = td
525 # Adjust Softmax to range from 0.0 to 1.0
526 elif node.op_type == "Softmax":
527 if not self.should_quantize_node(node):
528 continue
529 self.tensors_range[node.output[0]] = TensorData(lowest=np.float32(0.0), highest=np.float32(1.0))
530