latest (dev)
Copy
Latest development documentation · Updated 2026-10-15
Source code for tensorplay.ao.quantization.observer
"""Calibration statistics and uniform quantization parameter selection."""
import math
import warnings
from abc import ABC, abstractmethod
from collections import OrderedDict
from dataclasses import dataclass
from enum import Enum, auto
from functools import partial
import tensorplay
from tensorplay import nn
QUANT_MIN = -128
QUANT_MAX = 127
def is_per_tensor(qscheme):
return qscheme == tensorplay.per_tensor_affine or qscheme == tensorplay.per_tensor_symmetric
def is_per_channel(qscheme):
return qscheme in [tensorplay.per_channel_affine, tensorplay.per_channel_affine_float_qparams, tensorplay.per_channel_symmetric]
def check_min_max_valid(min_val, max_val):
if min_val.numel() == 0 or max_val.numel() == 0:
warnings.warn('must run observer before calling calculate_qparams. ' + 'Returning default values.', stacklevel=2)
return False
if min_val.dim() == 0 or max_val.dim() == 0:
if min_val == float('inf') and max_val == float('-inf'):
warnings.warn('must run observer before calling calculate_qparams. ' + 'Returning default values.', stacklevel=2)
return False
if min_val > max_val:
raise AssertionError(f'min {min_val} should be less than max {max_val}')
elif tensorplay.any(min_val > max_val):
raise AssertionError(f'min {min_val} should be less than max {max_val}')
return True
def calculate_qmin_qmax(quant_min, quant_max, has_customized_qrange, dtype, reduce_range):
if has_customized_qrange:
if dtype in [tensorplay.qint32, tensorplay.int32]:
initial_quant_min, initial_quant_max = (0, 2 ** 32 - 1)
else:
initial_quant_min, initial_quant_max = (0, 255)
custom_quant_min, custom_quant_max = (quant_min, quant_max)
if custom_quant_min is not None and custom_quant_max is not None:
initial_quant_min, initial_quant_max = (custom_quant_min, custom_quant_max)
qrange_len = initial_quant_max - initial_quant_min + 1
if dtype in [tensorplay.qint8, tensorplay.int8]:
if not 0 < qrange_len <= 256:
raise AssertionError('quantization range should be positive and not exceed the maximum bit range (=256).')
elif dtype in [tensorplay.qint32, tensorplay.int32]:
if not 0 < qrange_len <= 2 ** 32:
raise AssertionError('quantization range should be positive and not exceed the maximum bit range (=4294967296).')
if reduce_range:
quant_min, quant_max = (quant_min // 2, quant_max // 2)
elif dtype in [tensorplay.qint8, tensorplay.int8]:
if reduce_range:
quant_min, quant_max = (-64, 63)
else:
quant_min, quant_max = (-128, 127)
elif dtype in [tensorplay.quint8, tensorplay.uint8]:
if reduce_range:
quant_min, quant_max = (0, 127)
else:
quant_min, quant_max = (0, 255)
elif dtype in [tensorplay.qint32, tensorplay.int32]:
quant_min, quant_max = (-1 * 2 ** 31, 2 ** 31 - 1)
elif dtype == tensorplay.uint16:
quant_min, quant_max = (0, 2 ** 16 - 1)
elif dtype == tensorplay.int16:
quant_min, quant_max = (-2 ** 15, 2 ** 15 - 1)
else:
quant_min, quant_max = (0, 15)
return (quant_min, quant_max)
def validate_qmin_qmax(quant_min, quant_max):
if not quant_min <= 0 <= quant_max:
raise AssertionError('Used-specified quantization range must include 0.')
if quant_min >= quant_max:
raise AssertionError('qmin must be strictly less than qmax for user-specified quantization range.')
class _PartialWrapper:
def __init__(self, p):
self.p = p
self.callable_args = {}
def __call__(self, *args, **keywords):
for arg_name in self.callable_args:
if arg_name not in keywords:
keywords = {**keywords, arg_name: self.callable_args[arg_name]()}
return self.p(*args, **keywords)
def __repr__(self):
return self.p.__repr__() + self.callable_args.__repr__()
def with_args(self, **kwargs):
return _with_args(self, **kwargs)
def with_callable_args(self, **kwargs):
result = _PartialWrapper(p=self.p)
result.callable_args = {**self.callable_args, **kwargs}
return result
def _with_args(cls_or_self, **kwargs):
r = _PartialWrapper(partial(cls_or_self, **kwargs))
return r
def _with_callable_args(cls_or_self, **kwargs):
r = _PartialWrapper(partial(cls_or_self))
return r.with_callable_args(**kwargs)
class ObserverBase(nn.Module):
def __init__(self, dtype, is_dynamic=False):
super().__init__()
self.dtype = dtype
self.is_dynamic = is_dynamic
def forward(self, x):
pass
def calculate_qparams(self, **kwargs):
pass
with_args = classmethod(_with_args)
with_callable_args = classmethod(_with_callable_args)
class UniformQuantizationObserverBase(ObserverBase):
_version = 3
def __init__(self, dtype=tensorplay.quint8, qscheme=tensorplay.per_tensor_affine, reduce_range=False, quant_min=None, quant_max=None, factory_kwargs=None, eps=tensorplay.finfo(tensorplay.float32).eps, is_dynamic=False, **kwargs):
factory_kwargs = tensorplay.nn.factory_kwargs(factory_kwargs)
super().__init__(dtype=dtype, is_dynamic=is_dynamic, **kwargs)
self.qscheme = qscheme
if reduce_range:
warnings.warn('Please use quant_min and quant_max to specify the range for observers. reduce_range will be deprecated in a future release of quantization.', stacklevel=2)
self.reduce_range = reduce_range
self.register_buffer('eps', tensorplay.tensor([eps], **factory_kwargs))
if self.qscheme not in (tensorplay.per_tensor_affine, tensorplay.per_tensor_symmetric, tensorplay.per_channel_affine, tensorplay.per_channel_symmetric, tensorplay.per_channel_affine_float_qparams):
raise AssertionError('Default Observer only works for per_tensor_affine, per_tensor_symmetric, per_channel_affine, per_channel_symmetric and per_channel_float_qparams quantization scheme')
_ALLOWED_DTYPES = (tensorplay.qint8, tensorplay.quint8, tensorplay.quint4x2, tensorplay.quint1x8, tensorplay.qint32, tensorplay.int8, tensorplay.uint8, tensorplay.int16, tensorplay.int32, tensorplay.float8_e5m2, tensorplay.float8_e4m3fn, tensorplay.uint16)
if self.dtype not in _ALLOWED_DTYPES:
raise AssertionError(f'Default Observer only works for {_ALLOWED_DTYPES} data type')
self.has_customized_qrange = quant_min is not None and quant_max is not None
if self.has_customized_qrange:
validate_qmin_qmax(quant_min, quant_max)
self.quant_min, self.quant_max = calculate_qmin_qmax(quant_min, quant_max, self.has_customized_qrange, self.dtype, self.reduce_range)
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
version = local_metadata.get('version', None)
if version is None or version == 1:
eps = tensorplay.tensor([tensorplay.finfo(tensorplay.float32).eps])
state_dict[prefix + 'eps'] = eps
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
def _validate_qmin_qmax(self, quant_min, quant_max):
if not quant_min <= 0 <= quant_max:
raise AssertionError('Used-specified quantization range must include 0.')
if quant_min >= quant_max:
raise AssertionError('qmin must be strictly less than qmax for user-specified quantization range.')
def _derive_qparams(self, min_val, max_val):
if not check_min_max_valid(min_val, max_val):
return (tensorplay.tensor([1.0], device=min_val.device.type), tensorplay.tensor([0], device=min_val.device.type))
quant_min, quant_max = (self.quant_min, self.quant_max)
min_val_neg = tensorplay.minimum(min_val, tensorplay.zeros_like(min_val))
max_val_pos = tensorplay.maximum(max_val, tensorplay.zeros_like(max_val))
device = min_val_neg.device
scale = tensorplay.ones(min_val_neg.size(), dtype=tensorplay.float32, device=device)
zero_point = tensorplay.zeros(min_val_neg.size(), dtype=tensorplay.int64, device=device)
if self.qscheme == tensorplay.per_tensor_symmetric or self.qscheme == tensorplay.per_channel_symmetric:
max_val_pos = tensorplay.maximum(-min_val_neg, max_val_pos)
scale = _qparam_div(max_val_pos, float(quant_max - quant_min) / 2)
scale = tensorplay.maximum(scale, self.eps)
if self.dtype in [tensorplay.quint8, tensorplay.uint8]:
if self.has_customized_qrange:
zero_point = zero_point.new_full(zero_point.size(), (quant_min + quant_max) // 2)
else:
zero_point = zero_point.new_full(zero_point.size(), 128)
elif self.dtype == tensorplay.uint16:
zero_point = zero_point.new_full(zero_point.size(), 2 ** 15)
elif self.qscheme == tensorplay.per_channel_affine_float_qparams:
scale = _qparam_div(max_val - min_val, float(quant_max - quant_min))
scale = tensorplay.where(scale > self.eps, scale, tensorplay.ones_like(scale))
zero_point = -1 * min_val / scale
else:
scale = _qparam_div(max_val_pos - min_val_neg, float(quant_max - quant_min))
scale = tensorplay.maximum(scale, self.eps)
zero_point = quant_min - tensorplay.round(min_val_neg / scale).to(tensorplay.int)
zero_point = tensorplay.clamp(zero_point, quant_min, quant_max)
if scale.dim() == 0:
scale = scale.reshape(1)
if zero_point.dim() == 0:
zero_point = zero_point.reshape(1)
return (scale, zero_point)
def reset_min_max_vals(self):
raise NotImplementedError('Cannot reset min/max values in the given observer.')
[docs]
class MinMaxObserver(UniformQuantizationObserverBase):
def __init__(self, dtype=tensorplay.quint8, qscheme=tensorplay.per_tensor_affine, reduce_range=False, quant_min=None, quant_max=None, factory_kwargs=None, eps=tensorplay.finfo(tensorplay.float32).eps, is_dynamic=False, **kwargs):
if not is_per_tensor(qscheme):
raise NotImplementedError("MinMaxObserver's qscheme only support per_tensor_symmetric and per_tensor_affine.")
super().__init__(dtype=dtype, qscheme=qscheme, reduce_range=reduce_range, quant_min=quant_min, quant_max=quant_max, factory_kwargs=factory_kwargs, eps=eps, is_dynamic=is_dynamic, **kwargs)
factory_kwargs = tensorplay.nn.factory_kwargs(factory_kwargs)
self.register_buffer('min_val', tensorplay.tensor(float('inf'), **factory_kwargs))
self.register_buffer('max_val', tensorplay.tensor(float('-inf'), **factory_kwargs))
if self.qscheme == tensorplay.per_tensor_symmetric and self.reduce_range and (self.dtype == tensorplay.quint8):
raise NotImplementedError('Cannot reduce range for symmetric quantization for quint8')
def forward(self, x_orig):
if x_orig.numel() == 0:
return x_orig
x = x_orig.detach()
x = x.to(self.min_val.dtype)
min_val_cur, max_val_cur = tensorplay.aminmax(x)
min_val = tensorplay.minimum(min_val_cur, self.min_val)
max_val = tensorplay.maximum(max_val_cur, self.max_val)
self.min_val.copy_(min_val)
self.max_val.copy_(max_val)
return x_orig
def calculate_qparams(self):
return self._derive_qparams(self.min_val, self.max_val)
def extra_repr(self):
return f'min_val={self.min_val}, max_val={self.max_val}'
def reset_min_max_vals(self):
self.min_val.fill_(float('inf'))
self.max_val.fill_(float('-inf'))
[docs]
class MovingAverageMinMaxObserver(MinMaxObserver):
def __init__(self, averaging_constant=0.01, dtype=tensorplay.quint8, qscheme=tensorplay.per_tensor_affine, reduce_range=False, quant_min=None, quant_max=None, eps=tensorplay.finfo(tensorplay.float32).eps, is_dynamic=False, **kwargs):
if not is_per_tensor(qscheme):
raise NotImplementedError(f"MovingAverageMinMaxObserver's qscheme only support per_tensor_symmetric and per_tensor_affine. but got: {qscheme}")
self.averaging_constant = averaging_constant
if is_dynamic and self.averaging_constant != 1:
raise NotImplementedError(f"MovingAverageMinMaxObserver doesn't support dynamic quantization for averaging constant of {self.averaging_constant}")
super().__init__(dtype=dtype, qscheme=qscheme, reduce_range=reduce_range, quant_min=quant_min, quant_max=quant_max, eps=eps, is_dynamic=is_dynamic, **kwargs)
def forward(self, x_orig):
if x_orig.numel() == 0:
return x_orig
x = x_orig.detach()
x = x.to(self.min_val.dtype)
min_val = self.min_val
max_val = self.max_val
if min_val == float('inf') and max_val == float('-inf'):
min_val, max_val = tensorplay.aminmax(x)
else:
min_val_cur, max_val_cur = tensorplay.aminmax(x)
min_val = min_val + self.averaging_constant * (min_val_cur - min_val)
max_val = max_val + self.averaging_constant * (max_val_cur - max_val)
self.min_val.copy_(min_val)
self.max_val.copy_(max_val)
return x_orig
[docs]
class PerChannelMinMaxObserver(UniformQuantizationObserverBase):
def __init__(self, ch_axis=0, dtype=tensorplay.quint8, qscheme=tensorplay.per_channel_affine, reduce_range=False, quant_min=None, quant_max=None, factory_kwargs=None, eps=tensorplay.finfo(tensorplay.float32).eps, is_dynamic=False, **kwargs):
if not is_per_channel(qscheme):
raise NotImplementedError("PerChannelMinMaxObserver's qscheme only support per_channel_symmetric, per_channel_affine and per_channel_affine_float_qparams.")
if is_dynamic:
raise NotImplementedError("PerChannelMinMaxObserver doesn't support dynamic quantization")
super().__init__(dtype=dtype, qscheme=qscheme, reduce_range=reduce_range, quant_min=quant_min, quant_max=quant_max, factory_kwargs=factory_kwargs, eps=eps, is_dynamic=is_dynamic, **kwargs)
factory_kwargs = tensorplay.nn.factory_kwargs(factory_kwargs)
self.ch_axis = ch_axis
self.register_buffer('min_val', tensorplay.tensor([], **factory_kwargs))
self.register_buffer('max_val', tensorplay.tensor([], **factory_kwargs))
if self.qscheme == tensorplay.per_channel_symmetric and self.reduce_range and (self.dtype == tensorplay.quint8):
raise NotImplementedError('Cannot reduce range for symmetric quantization for quint8')
def forward(self, x_orig):
return self._forward(x_orig)
def _forward(self, x_orig):
if x_orig.numel() == 0:
return x_orig
x = x_orig.detach()
min_val = self.min_val
max_val = self.max_val
x_dim = x.size()
new_axis_list = [i for i in range(len(x_dim))]
new_axis_list[self.ch_axis] = 0
new_axis_list[0] = self.ch_axis
y = x.permute(new_axis_list)
y = y.to(self.min_val.dtype)
y = tensorplay.flatten(y, start_dim=1)
if min_val.numel() == 0 or max_val.numel() == 0:
min_val, max_val = tensorplay.aminmax(y, dim=1)
else:
min_val_cur, max_val_cur = tensorplay.aminmax(y, dim=1)
min_val = tensorplay.minimum(min_val_cur, min_val)
max_val = tensorplay.maximum(max_val_cur, max_val)
self.min_val.resize_(min_val.shape)
self.max_val.resize_(max_val.shape)
self.min_val.copy_(min_val)
self.max_val.copy_(max_val)
return x_orig
def calculate_qparams(self):
return self._derive_qparams(self.min_val, self.max_val)
def extra_repr(self):
return f'min_val={self.min_val}, max_val={self.max_val}'
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs):
for name in ('min_val', 'max_val'):
old_name = name.replace('_val', '_vals')
if prefix + old_name in state_dict and prefix + name not in state_dict:
state_dict[prefix + name] = state_dict.pop(prefix + old_name)
if prefix + name in state_dict:
getattr(self, name).resize_(state_dict[prefix + name].shape)
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs)
def reset_min_max_vals(self):
self.min_val.resize_(0)
self.max_val.resize_(0)
[docs]
class MovingAveragePerChannelMinMaxObserver(PerChannelMinMaxObserver):
def __init__(self, averaging_constant=0.01, ch_axis=0, dtype=tensorplay.quint8, qscheme=tensorplay.per_channel_affine, reduce_range=False, quant_min=None, quant_max=None, eps=tensorplay.finfo(tensorplay.float32).eps, is_dynamic=False, **kwargs):
if not is_per_channel(qscheme):
raise NotImplementedError("MovingAveragePerChannelMinMaxObserver's qscheme only support per_channel_symmetric, per_channel_affine and per_channel_affine_float_qparams.")
if is_dynamic:
raise NotImplementedError("MovingAveragePerChannelMinMaxObserver doesn't support dynamic quantization")
super().__init__(ch_axis=ch_axis, dtype=dtype, qscheme=qscheme, reduce_range=reduce_range, quant_min=quant_min, quant_max=quant_max, eps=eps, is_dynamic=is_dynamic, **kwargs)
self.averaging_constant = averaging_constant
def forward(self, x_orig):
if x_orig.numel() == 0:
return x_orig
x = x_orig.detach()
x = x.to(self.min_val.dtype)
min_val = self.min_val
max_val = self.max_val
x_dim = x.size()
new_axis_list = [i for i in range(len(x_dim))]
new_axis_list[self.ch_axis] = 0
new_axis_list[0] = self.ch_axis
y = x.permute(new_axis_list)
y = tensorplay.flatten(y, start_dim=1)
if min_val.numel() == 0 or max_val.numel() == 0:
min_val, max_val = tensorplay.aminmax(y, dim=1)
else:
min_val_cur, max_val_cur = tensorplay.aminmax(y, dim=1)
min_val = min_val + self.averaging_constant * (min_val_cur - min_val)
max_val = max_val + self.averaging_constant * (max_val_cur - max_val)
self.min_val.resize_(min_val.shape)
self.max_val.resize_(max_val.shape)
self.min_val.copy_(min_val)
self.max_val.copy_(max_val)
return x_orig
[docs]
class HistogramObserver(UniformQuantizationObserverBase):
def __init__(self, bins=2048, dtype=tensorplay.quint8, qscheme=tensorplay.per_tensor_affine, reduce_range=False, quant_min=None, quant_max=None, factory_kwargs=None, eps=tensorplay.finfo(tensorplay.float32).eps, is_dynamic=False, **kwargs):
if not is_per_tensor(qscheme):
raise NotImplementedError("HistogramObserver's qscheme only support per_tensor_symmetric and per_tensor_affine.")
if is_dynamic:
raise NotImplementedError("HistogramObserver doesn't support dynamic quantization")
super().__init__(dtype=dtype, qscheme=qscheme, reduce_range=reduce_range, quant_min=quant_min, quant_max=quant_max, factory_kwargs=factory_kwargs, eps=eps, is_dynamic=is_dynamic, **kwargs)
factory_kwargs = tensorplay.nn.factory_kwargs(factory_kwargs)
self.bins = bins
self.register_buffer('histogram', tensorplay.zeros(self.bins, **factory_kwargs))
self.register_buffer('min_val', tensorplay.tensor(float('inf'), **factory_kwargs))
self.register_buffer('max_val', tensorplay.tensor(float('-inf'), **factory_kwargs))
self.dst_nbins = 2 ** (32 if self.dtype in (tensorplay.qint32, tensorplay.int32) else 8 if self.dtype in (tensorplay.qint8, tensorplay.quint8) else tensorplay.iinfo(self.dtype).bits)
self.upsample_rate = 16
def _get_norm(self, delta_begin, delta_end, density):
norm = (delta_end * delta_end * delta_end - delta_begin * delta_begin * delta_begin) / 3
return density * norm
def _compute_quantization_error(self, next_start_bin, next_end_bin):
bin_width = (self.max_val.item() - self.min_val.item()) / self.bins
dst_bin_width = bin_width * (next_end_bin - next_start_bin + 1) / self.dst_nbins
if dst_bin_width == 0.0:
return 0.0
src_bin = tensorplay.arange(self.bins, device=self.histogram.device)
src_bin_begin = (src_bin - next_start_bin) * bin_width
src_bin_end = src_bin_begin + bin_width
dst_bin_of_begin = tensorplay.clamp(tensorplay.div(src_bin_begin, dst_bin_width, rounding_mode='floor'), 0, self.dst_nbins - 1)
dst_bin_of_begin_center = (dst_bin_of_begin + 0.5) * dst_bin_width
dst_bin_of_end = tensorplay.clamp(tensorplay.div(src_bin_end, dst_bin_width, rounding_mode='floor'), 0, self.dst_nbins - 1)
density = self.histogram / bin_width
norm = tensorplay.zeros(self.bins, device=self.histogram.device)
delta_begin = src_bin_begin - dst_bin_of_begin_center
delta_end = dst_bin_width / 2
norm += self._get_norm(delta_begin, tensorplay.ones(self.bins, device=self.histogram.device) * delta_end, density)
norm += (dst_bin_of_end - dst_bin_of_begin - 1) * self._get_norm(-dst_bin_width / 2, dst_bin_width / 2, density)
dst_bin_of_end_center = dst_bin_of_end * dst_bin_width + dst_bin_width / 2
delta_begin = -dst_bin_width / 2
delta_end = src_bin_end - dst_bin_of_end_center
norm += self._get_norm(delta_begin, delta_end, density)
return norm.sum().item()
def _non_linear_param_search(self):
if self.histogram.size()[0] != self.bins:
raise AssertionError('bins mismatch')
bin_width = (self.max_val - self.min_val) / self.bins
total = tensorplay.sum(self.histogram).item()
cSum = tensorplay.cumsum(self.histogram, dim=0)
stepsize = 1e-05
alpha = 0.0
beta = 1.0
start_bin = 0
end_bin = self.bins - 1
norm_min = float('inf')
while alpha < beta:
next_alpha = alpha + stepsize
next_beta = beta - stepsize
l = start_bin
r = end_bin
while l < end_bin and cSum[l] < next_alpha * total:
l = l + 1
while r > start_bin and cSum[r] > next_beta * total:
r = r - 1
next_start_bin = start_bin
next_end_bin = end_bin
if l - start_bin > end_bin - r:
next_start_bin = l
alpha = next_alpha
else:
next_end_bin = r
beta = next_beta
if next_start_bin == start_bin and next_end_bin == end_bin:
continue
norm = self._compute_quantization_error(next_start_bin, next_end_bin)
if norm > norm_min:
break
norm_min = norm
start_bin = next_start_bin
end_bin = next_end_bin
new_min = self.min_val + bin_width * start_bin
new_max = self.min_val + bin_width * (end_bin + 1)
return (new_min, new_max)
def _upscale_histogram(self, histogram, orig_min, orig_max, update_min, update_max):
histogram = histogram.repeat_interleave(self.upsample_rate) / self.upsample_rate
bin_size = (orig_max - orig_min) / (self.bins * self.upsample_rate)
mid_points_histogram = tensorplay.linspace(float(orig_min), float(orig_max), self.bins * self.upsample_rate + 1, device=orig_min.device)[:-1].to(histogram.device) + 0.5 * bin_size
boundaries_new_histogram = tensorplay.linspace(float(update_min), float(update_max), self.bins + 1, device=update_min.device).to(histogram.device)
bucket_assignments = tensorplay.bucketize(mid_points_histogram, boundaries_new_histogram, right=True) - 1
bucket_assignments[bucket_assignments >= self.bins] = self.bins - 1
bucket_assignments[bucket_assignments < 0] = 0
update_histogram = tensorplay.bincount(bucket_assignments, weights=histogram, minlength=self.bins)
return update_histogram
def _combine_histograms(self, orig_hist, orig_min, orig_max, update_hist, update_min, update_max):
if update_min == orig_min and update_max == orig_max:
return orig_hist + update_hist
if orig_min == orig_max:
bin_value = tensorplay.sum(orig_hist)
transformed_orig_hist = tensorplay.histc(orig_min, bins=self.bins, min=float(update_min), max=float(update_max)) * bin_value
return transformed_orig_hist + update_hist
if update_min > orig_min:
raise AssertionError('update_min must be <= orig_min')
if update_max < orig_max:
raise AssertionError('update_max must be >= orig_max')
transformed_orig_hist = self._upscale_histogram(orig_hist, orig_min, orig_max, update_min, update_max)
return update_hist + transformed_orig_hist
def reset_histogram(self, x, min_val, max_val):
self.min_val.resize_(min_val.shape)
self.min_val.copy_(min_val)
self.max_val.resize_(max_val.shape)
self.max_val.copy_(max_val)
if min_val.numel() != 1 or max_val.numel() != 1:
raise AssertionError('histogram min/max values must be scalar.')
new_histogram = tensorplay.histc(x, self.bins, min=float(min_val), max=float(max_val))
self.histogram.resize_(new_histogram.shape)
self.histogram.copy_(new_histogram)
def forward(self, x_orig):
if x_orig.numel() == 0:
return x_orig
x = x_orig.detach()
x_min, x_max = tensorplay.aminmax(x)
if x_min == -tensorplay.inf or x_max == tensorplay.inf:
warnings.warn('inf detected in input tensor, ignoring input', stacklevel=2)
x = x[x.abs() != tensorplay.inf]
if x.numel() == 0:
return x_orig
x_min, x_max = tensorplay.aminmax(x)
current_min = self.min_val
current_max = self.max_val
is_uninitialized = self.min_val == float('inf') or self.max_val == float('-inf')
if is_uninitialized:
self.reset_histogram(x, x_min, x_max)
else:
update_min, update_max = (x_min, x_max)
new_min = tensorplay.minimum(current_min, update_min)
new_max = tensorplay.maximum(current_max, update_max)
new_min, new_max = (new_min.detach(), new_max.detach())
update_histogram = tensorplay.histc(x, self.bins, min=float(new_min), max=float(new_max)).to(self.histogram.device)
if new_min == current_min and new_max == current_max:
combined_histogram = self.histogram + update_histogram
self.histogram.resize_(combined_histogram.shape)
self.histogram.copy_(combined_histogram)
else:
combined_histogram = self._combine_histograms(self.histogram, current_min, current_max, update_histogram, new_min, new_max)
self.histogram.resize_(combined_histogram.shape)
self.histogram.copy_(combined_histogram)
self.min_val.resize_(new_min.shape)
self.min_val.copy_(new_min)
self.max_val.resize_(new_max.shape)
self.max_val.copy_(new_max)
return x_orig
def calculate_qparams(self):
is_uninitialized = self.min_val == float('inf') and self.max_val == float('-inf')
if is_uninitialized:
warnings.warn('must run observer before calling calculate_qparams. Returning default scale and zero point ', stacklevel=2)
return (tensorplay.tensor([1.0], device=self.min_val.device.type), tensorplay.tensor([0], device=self.min_val.device.type))
if self.bins != len(self.histogram):
raise AssertionError('The number of bins in histogram should be equal to the number of bins supplied while making this observer')
new_min, new_max = self._non_linear_param_search()
return self._derive_qparams(new_min, new_max)
def _save_to_state_dict(self, destination, prefix, keep_vars):
super()._save_to_state_dict(destination, prefix, keep_vars)
destination[prefix + 'min_val'] = self.min_val
destination[prefix + 'max_val'] = self.max_val
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
version = local_metadata.get('version', None)
if version is None or version < 3:
min_val_name, max_val_name = (prefix + 'min_val', prefix + 'max_val')
if min_val_name in state_dict:
if state_dict[min_val_name].shape == tensorplay.Size([0]):
state_dict[min_val_name] = tensorplay.tensor(float('inf'))
if max_val_name in state_dict:
if state_dict[max_val_name].shape == tensorplay.Size([0]):
state_dict[max_val_name] = tensorplay.tensor(float('-inf'))
local_state = ['min_val', 'max_val']
for name in local_state:
key = prefix + name
if key in state_dict:
val = state_dict[key]
setattr(self, name, val)
elif strict:
missing_keys.append(key)
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
def extra_repr(self):
return f'min_val={self.min_val}, max_val={self.max_val}'
[docs]
class FixedQParamsObserver(ObserverBase):
def __init__(self, scale, zero_point, dtype=tensorplay.quint8, qscheme=tensorplay.per_tensor_affine, quant_min=0, quant_max=255, is_dynamic=False, **kwargs):
if is_dynamic:
raise NotImplementedError("FixedQParamsObserver doesn't support dynamic quantization")
super().__init__(dtype=dtype, is_dynamic=is_dynamic, **kwargs)
self.quant_min = quant_min
self.quant_max = quant_max
self.register_buffer('scale', tensorplay.tensor([scale], dtype=tensorplay.float))
self.register_buffer('zero_point', tensorplay.tensor([zero_point], dtype=tensorplay.int))
self.dtype = dtype
self.qscheme = qscheme
def forward(self, X):
return X
def calculate_qparams(self):
return (self.scale, self.zero_point)
[docs]
class PlaceholderObserver(ObserverBase):
def __init__(self, dtype=tensorplay.float32, custom_op_name='', compute_dtype=None, quant_min=None, quant_max=None, qscheme=None, eps=None, is_dynamic=False):
super().__init__(dtype=dtype, is_dynamic=is_dynamic)
if qscheme is None:
qscheme = tensorplay.per_tensor_affine
if eps is None:
eps = tensorplay.finfo(tensorplay.float32).eps
self.dtype = dtype
self.qscheme = qscheme
self.quant_min = quant_min
self.quant_max = quant_max
self.eps = eps
self.custom_op = custom_op_name
if compute_dtype:
is_dynamic = True
warnings.warn('Please use `is_dynamic` instead of `compute_dtype`. `compute_dtype` will be deprecated in a future release of quantization.', stacklevel=2)
def forward(self, x):
return x
def extra_repr(self):
return f'dtype={self.dtype}, is_dynamic={self.is_dynamic}'
def calculate_qparams(self):
raise Exception('calculate_qparams should not be called for PlaceholderObserver')
class RecordingObserver(ObserverBase):
__annotations__ = {'tensor_val': list[tensorplay.Tensor | None]}
def __init__(self, dtype=tensorplay.quint8):
super().__init__(dtype=dtype, is_dynamic=False)
self.tensor_val = []
def forward(self, x):
self.tensor_val.append(x.clone())
return x
def calculate_qparams(self):
raise Exception('calculate_qparams should not be called for RecordingObserver')
def get_tensor_value(self):
return self.tensor_val
class NoopObserver(ObserverBase):
def __init__(self, dtype=tensorplay.float16, custom_op_name=''):
super().__init__(dtype=dtype, is_dynamic=False)
self.dtype = dtype
self.custom_op = custom_op_name
def forward(self, x):
return x
def calculate_qparams(self):
raise Exception('calculate_qparams should not be called for NoopObserver')
class ReuseInputObserver(ObserverBase):
def __init__(self):
super().__init__(tensorplay.quint8, is_dynamic=False)
def forward(self, x):
return x
def calculate_qparams(self):
raise Exception('calculate_qparams should not be called for ReuseInputObserver')
class MappingType(Enum):
"""How a floating point range is mapped onto the integer grid.
SYMMETRIC maps the symmetric float range around zero onto the full
integer range. SYMMETRIC_NO_CLIPPING_ERR derives the scale from the
per-side min/max separately so negative values round with less error.
ASYMMETRIC maps the observed float range directly onto the integer range.
"""
SYMMETRIC = auto()
SYMMETRIC_NO_CLIPPING_ERR = auto()
ASYMMETRIC = auto()
class ZeroPointDomain(Enum):
"""Whether the zero point lives in the integer or floating point domain.
INT: quantized = round(float / scale) + zero_point (integer)
FLOAT: quantized = round((float - (zero_point - scale * mid)) / scale)
NONE: quantized = round(float / scale)
"""
INT = auto()
FLOAT = auto()
NONE = auto()
class TorchAODType(Enum):
"""Placeholder for narrow integer dtypes without a native representation."""
INT1 = auto()
INT2 = auto()
INT3 = auto()
INT4 = auto()
INT5 = auto()
INT6 = auto()
INT7 = auto()
@dataclass(frozen=True)
class Granularity:
"""Base class for quantization granularity descriptors."""
@dataclass(frozen=True)
class PerBlock(Granularity):
"""Per-block granularity with an explicit block size."""
block_size: tuple[int, ...]
@dataclass(frozen=True)
class PerTensor(Granularity):
"""A single set of quantization parameters for the whole tensor."""
@dataclass(frozen=True)
class PerAxis(Granularity):
"""One set of quantization parameters per slice along ``axis``."""
axis: int
@dataclass(frozen=True)
class PerGroup(Granularity):
"""One set of quantization parameters per group of ``group_size`` elements."""
group_size: int
class PerRow(Granularity):
"""Row-wise granularity, a per-axis special case used by float8 matmuls."""
class PerToken(Granularity):
"""One set of quantization parameters per token (last dimension slice)."""
def get_block_size(
input_shape: tuple[int, ...], granularity: Granularity
) -> tuple[int, ...]:
"""Derive the block size from an input shape and a granularity type."""
if not isinstance(granularity, Granularity):
raise AssertionError(
"Please provide an instance of Granularity, not subclass of it"
)
if isinstance(granularity, PerTensor):
return input_shape
elif isinstance(granularity, PerAxis):
block_size = list(input_shape)
block_size[granularity.axis] = 1
return tuple(block_size)
elif isinstance(granularity, PerRow):
return (1,) * (len(input_shape) - 1) + (input_shape[-1],)
elif isinstance(granularity, PerGroup):
if len(input_shape) != 2:
raise AssertionError(
f"Expecting input shape dim to be 2 for per group quantization, gotinput shape: {input_shape}"
)
return (1, granularity.group_size)
elif isinstance(granularity, PerToken):
block_size = [1] * len(input_shape)
block_size[-1] = input_shape[-1]
return tuple(block_size)
raise ValueError(f"Unsupported Granularity: {granularity}")
class AffineQuantizedObserverBase(ABC, nn.Module):
"""Observer for affine quantization with an explicit granularity.
Stores the mapping configuration plus the scale/zero-point dtypes and
zero-point domain used by affine quantized tensors.
"""
with_args = classmethod(_with_args)
def __init__(
self,
mapping_type: MappingType,
target_dtype: tensorplay.dtype,
granularity: Granularity,
quant_min: int | None = None,
quant_max: int | None = None,
eps: float | None = None,
scale_dtype: tensorplay.dtype | None = None,
zero_point_dtype: tensorplay.dtype | None = None,
preserve_zero: bool = True,
zero_point_domain: ZeroPointDomain | None = ZeroPointDomain.INT,
**kwargs,
):
super().__init__()
if granularity is None:
raise AssertionError("granularity is None")
self.mapping_type = mapping_type
self.target_dtype = target_dtype
self.granularity = granularity
self.quant_min = quant_min
self.quant_max = quant_max
self.eps = eps
self.scale_dtype = scale_dtype
self.zero_point_dtype = zero_point_dtype
self.preserve_zero = preserve_zero
self.zero_point_domain = zero_point_domain
self.block_size = None
self.original_dtype = None
@abstractmethod
def forward(self, input):
"""Record statistics and return the original input tensor."""
@abstractmethod
def calculate_qparams(self) -> tuple:
"""Return the (scale, zero_point) tuple for this observer."""
def convert(self, model, observer_node):
"""Rewrite the observer node into its quantized representation.
Requires the graph-mode lowering hooks for affine quantized tensors,
which are not part of this package yet.
"""
raise NotImplementedError(
"AffineQuantizedObserverBase.convert requires graph-mode affine lowering"
)
def with_args(**kwargs):
return lambda cls: _with_args(cls, **kwargs)
def _record(self, x):
return self(x)
def _observation_state(self):
return self.state_dict()
def _load_observation_state(self, state):
return self.load_state_dict(state)
def _legacy_qparams(min_val, max_val, quant_min=QUANT_MIN, quant_max=QUANT_MAX):
lo, hi = min(float(min_val), 0.), max(float(max_val), 0.)
scale = max((hi - lo) / (quant_max - quant_min), 1.1920928955078125e-7)
return scale, max(quant_min, min(quant_max, quant_min - round(lo / scale)))
ObserverBase.record = _record
ObserverBase.observation_state = _observation_state
ObserverBase.load_observation_state = _load_observation_state
ObserverBase._calculate_qparams = staticmethod(_legacy_qparams)
MinMaxObserver.reset = MinMaxObserver.reset_min_max_vals
PerChannelMinMaxObserver.reset = PerChannelMinMaxObserver.reset_min_max_vals
def _reset_histogram(self):
self.histogram.zero_()
self.min_val.fill_(float('inf'))
self.max_val.fill_(float('-inf'))
HistogramObserver.reset = _reset_histogram
default_observer = MinMaxObserver.with_args(quant_min=0, quant_max=127)
default_weight_observer = MinMaxObserver.with_args(dtype=tensorplay.qint8, qscheme=tensorplay.per_tensor_symmetric)
default_per_channel_weight_observer = PerChannelMinMaxObserver.with_args(dtype=tensorplay.qint8, qscheme=tensorplay.per_channel_symmetric)
default_histogram_observer = HistogramObserver.with_args(quant_min=0, quant_max=127)
default_dynamic_quant_observer = PlaceholderObserver.with_args(dtype=tensorplay.quint8, quant_min=0, quant_max=255, is_dynamic=True)
default_float_qparams_observer = PerChannelMinMaxObserver.with_args(dtype=tensorplay.quint8, qscheme=tensorplay.per_channel_affine_float_qparams)
default_float_qparams_observer_4bit = PerChannelMinMaxObserver.with_args(dtype=tensorplay.quint4x2, qscheme=tensorplay.per_channel_affine_float_qparams,ch_axis=0)
default_fixed_qparams_range_neg1to1_observer = FixedQParamsObserver.with_args(scale=2./256, zero_point=128)
default_fixed_qparams_range_0to1_observer = FixedQParamsObserver.with_args(scale=1./256, zero_point=0)
default_placeholder_observer = PlaceholderObserver
weight_observer_range_neg_127_to_127 = MinMaxObserver.with_args(
dtype=tensorplay.qint8,
qscheme=tensorplay.per_tensor_symmetric,
quant_min=-127,
quant_max=127,
eps=2**-12,
)
per_channel_weight_observer_range_neg_127_to_127 = PerChannelMinMaxObserver.with_args(
dtype=tensorplay.qint8,
qscheme=tensorplay.per_channel_symmetric,
quant_min=-127,
quant_max=127,
eps=2**-12,
)
default_debug_observer = RecordingObserver
default_reuse_input_observer = ReuseInputObserver
default_symmetric_fixed_qparams_observer = default_fixed_qparams_range_neg1to1_observer
default_affine_fixed_qparams_observer = default_fixed_qparams_range_0to1_observer
[docs]
def get_observer_state_dict(model):
from .fake_quantize import FakeQuantizeBase
state = OrderedDict()
for name, module in model.named_modules():
if isinstance(module, (ObserverBase, FakeQuantizeBase)):
module._save_to_state_dict(state, name + '.' if name else '', False)
return state
[docs]
def load_observer_state_dict(model, obs_dict):
expected = get_observer_state_dict(model)
if set(expected) != set(obs_dict):
raise RuntimeError('Observer state keys do not match the model')
state = model.state_dict()
state.update(obs_dict)
model.load_state_dict(state)
def _qparam_div(value, denominator):
if value.device.type == 'cuda':
return value * (1.0 / denominator)
return value / denominatorHelp improve this page
Found an error, an unclear step, or a missing example?
Was this page helpful?