Feat (a2q+): improving accumulator-aware weight quantization (#797)

src/brevitas/core/scaling/__init__.py

@@ -7,6 +7,7 @@
 from .int_scaling import IntScaling
 from .int_scaling import PowerOfTwoIntScaling
 from .pre_scaling import AccumulatorAwareParameterPreScaling
+from .pre_scaling import AccumulatorAwareZeroCenterParameterPreScaling
 from .pre_scaling import ParameterPreScalingWeightNorm
 from .runtime import RuntimeStatsScaling
 from .runtime import StatsFromParameterScaling

src/brevitas/core/scaling/pre_scaling.py

@@ -13,10 +13,14 @@
 from brevitas.core.restrict_val import _RestrictClampValue
 from brevitas.core.stats import SCALAR_SHAPE
 from brevitas.core.stats.stats_wrapper import _Stats
+from brevitas.core.zero_point import PreZeroCenterZeroPoint
 from brevitas.function import abs_binary_sign_grad
 from brevitas.function import get_upper_bound_on_l1_norm
 
-__all__ = ["ParameterPreScalingWeightNorm", "AccumulatorAwareParameterPreScaling"]
+__all__ = [
+    "ParameterPreScalingWeightNorm",
+    "AccumulatorAwareParameterPreScaling",
+    "AccumulatorAwareZeroCenterParameterPreScaling"]
 
 
 class ParameterPreScalingWeightNorm(brevitas.jit.ScriptModule):
@@ -113,7 +117,7 @@ def _load_from_state_dict(
 class AccumulatorAwareParameterPreScaling(ParameterPreScalingWeightNorm):
     """
     ScriptModule implementation of learned pre-clipping scaling factor to support
-    accumulator-aware quantizaion (A2Q) as proposed in `A2Q: Accumulator-Aware Quantization
+    accumulator-aware quantization (A2Q) as proposed in `A2Q: Accumulator-Aware Quantization
     with Guaranteed Overflow Avoidance` by I. Colbert, A. Pappalardo, and J. Petri-Koenig.
 
     The module parameterizes the pre-clipping scaling factor (i.e., `pre_scale`) of the
@@ -150,38 +154,118 @@ class AccumulatorAwareParameterPreScaling(ParameterPreScalingWeightNorm):
     """
 
     def __init__(
-        self,
-        scaling_impl: Module,
-        normalize_stats_impl: Module,
-        accumulator_bit_width_impl: Module,
-        scaling_stats_input_view_shape_impl: Module,
-        tracked_parameter_list: List[torch.nn.Parameter],
-        pre_scaling_shape: Optional[Tuple[int, ...]] = None,
-        restrict_pre_scaling_impl: Optional[Module] = None,
-        pre_scaling_min_val: Optional[float] = None,
-    ) -> None:
+            self,
+            scaling_impl: Module,
+            normalize_stats_impl: Module,
+            accumulator_bit_width_impl: Module,
+            scaling_stats_input_view_shape_impl: Module,
+            tracked_parameter_list: List[torch.nn.Parameter],
+            pre_scaling_shape: Optional[Tuple[int, ...]] = None,
+            restrict_pre_scaling_impl: Optional[Module] = None,
+            pre_scaling_min_val: Optional[float] = None) -> None:
         super().__init__(
             scaling_impl,
             normalize_stats_impl,
             scaling_stats_input_view_shape_impl,
             tracked_parameter_list,
             pre_scaling_shape,
             restrict_pre_scaling_impl,
-            pre_scaling_min_val,
-        )
+            pre_scaling_min_val)
         self.accumulator_bit_width = accumulator_bit_width_impl
 
     @brevitas.jit.script_method
-    def forward(self, weights: Tensor, input_bit_width: Tensor, input_is_signed: bool) -> Tensor:
-        """Takes weights as input and returns the pre-clipping scaling factor"""
+    def calc_max_l1_norm(self, input_bit_width: Tensor, input_is_signed: bool) -> Tensor:
+        accumulator_bit_width = self.accumulator_bit_width()
+        upper_bound = get_upper_bound_on_l1_norm(
+            accumulator_bit_width, input_bit_width, input_is_signed)
+        return upper_bound
+
+    @brevitas.jit.script_method
+    def inner_forward(self, weights: Tensor, input_bit_width: Tensor, input_is_signed: bool):
         weights = self.stats_input_view_shape_impl(weights)
         d_w = self.stats(weights)  # denominator for weight normalization
         s = self.scaling_impl(weights)  # s
         g = abs_binary_sign_grad(self.restrict_clamp_scaling(self.value))  # g
-        T = get_upper_bound_on_l1_norm(
-            self.accumulator_bit_width(), input_bit_width, input_is_signed)  # T / s
+        T = self.calc_max_l1_norm(input_bit_width, input_is_signed)  # T / s
         g = torch.clamp_max(g / s, T)
         value = d_w / g  # calculating final pre-clipping scaling factor
         # re-apply clamp_min_ste from restrict_scaling_impl to the specified pre_scaling_min_val
         value = self.restrict_clamp_scaling.clamp_min_ste(value)
         return value
+
+    @brevitas.jit.script_method
+    def forward(self, weights: Tensor, input_bit_width: Tensor, input_is_signed: bool) -> Tensor:
+        """Takes weights, input bit-width, and input sign as input and returns the pre-clipping
+        scaling factor per-channel, which is $s \cdot \Vert v \Vert_1 / g$"""
+        value = self.inner_forward(weights, input_bit_width, input_is_signed)
+        return value
+
+
+class AccumulatorAwareZeroCenterParameterPreScaling(AccumulatorAwareParameterPreScaling):
+    """
+    ScriptModule implementation of learned pre-clipping scaling factor to support
+    A2Q+ as proposed in `A2Q+: Improving Accumulator-Aware Weight Quantization`.
+
+    The module implements the zero-centering constraint as a pre-clipping zero-point
+    (i.e., `PreZeroCenterZeroPoint`) to relax the l1-norm constraint.
+
+    Args:
+        scaling_impl (Module): post-clipping scaling factor.
+        pre_zero_point_impl (Module): pre-clipping zero-point.
+        normalize_stats_impl (Module): calculate statistics for normalizing weight parameter.
+        accumulator_bit_width_impl (Module): module that returns the accumulator bit-width.
+        scaling_stats_input_view_shape_impl (Module): transforming scaling to a new shape.
+        tracked_parameter_list (List[torch.nn.Parameter]): list of tracked weight parameters
+            for tensor quantizer.
+        pre_scaling_shape (Tuple[int]): shape of pre-clipping scaling factor. Default: None.
+        restrict_pre_scaling_impl (Module): restrict pre_scaling_init according to some
+            criteria. Default: None.
+        pre_scaling_min_val (float): force a lower-bound on scaling_init. Default: None.
+
+    Returns:
+        Tensor: scaling factor wrapped in a float torch.Tensor.
+    """
+
+    def __init__(
+            self,
+            scaling_impl: Module,
+            pre_zero_point_impl: Module,
+            normalize_stats_impl: Module,
+            accumulator_bit_width_impl: Module,
+            scaling_stats_input_view_shape_impl: Module,
+            tracked_parameter_list: List[Parameter],
+            pre_scaling_shape: Optional[Tuple[int, ...]] = None,
+            restrict_pre_scaling_impl: Optional[Module] = None,
+            pre_scaling_min_val: Optional[float] = None) -> None:
+        super().__init__(
+            scaling_impl,
+            normalize_stats_impl,
+            accumulator_bit_width_impl,
+            scaling_stats_input_view_shape_impl,
+            tracked_parameter_list,
+            pre_scaling_shape,
+            restrict_pre_scaling_impl,
+            pre_scaling_min_val)
+        assert isinstance(
+            pre_zero_point_impl, PreZeroCenterZeroPoint
+        ), "Error: A2Q+ requires a pre-clipping zero-centering zero-point."
+        self.pre_zero_point = pre_zero_point_impl
+
+    @brevitas.jit.script_method
+    def calc_max_l1_norm(self, input_bit_width: Tensor, input_is_signed: bool) -> Tensor:
+        """ """
+        assert input_bit_width is not None, "A2Q+ relies on input bit-width."
+        max_accumulator_bit_width = self.accumulator_bit_width()  # P
+        max_accumulator_mag = pow(2.0, max_accumulator_bit_width) - 2.0  # 2^P - 2
+        max_input_mag = pow(2.0, input_bit_width) - 1.0  # 2^N - 1
+        return max_accumulator_mag / max_input_mag
+
+    @brevitas.jit.script_method
+    def forward(self, weights: Tensor, input_bit_width: Tensor, input_is_signed: bool) -> Tensor:
+        """Takes weights, input bit-width, and input sign as input and returns the pre-clipping
+        scaling factor per-channel, which is $s \cdot \Vert v - \mu_v \Vert_1 / g$"""
+        # NOTE: A2Q+ requires zero-centering the floating-point weights, which means that the
+        # calculation of the l1-norm needs to be done over the zero-centered weights.
+        z = self.pre_zero_point.get_zero_center(weights)
+        value = self.inner_forward(weights + z, input_bit_width, input_is_signed)
+        return value

src/brevitas/core/zero_point.py

@@ -24,7 +24,8 @@
     'StatsFromParameterZeroPoint',
     'ParameterFromRuntimeZeroPoint',
     'ParameterZeroPoint',
-    'ParameterFromStatsFromParameterZeroPoint']
+    'ParameterFromStatsFromParameterZeroPoint',
+    'PreZeroCenterZeroPoint']
 
 
 class ZeroZeroPoint(brevitas.jit.ScriptModule):
@@ -294,3 +295,33 @@ def _load_from_state_dict(
             self.init_done = True
         if config.IGNORE_MISSING_KEYS and value_key in missing_keys:
             missing_keys.remove(value_key)
+
+
+class PreZeroCenterZeroPoint(brevitas.jit.ScriptModule):
+    """Experimental ScriptModule implementation of a pre-scaling zero-point that zero-centers
+    the incoming tensors. This is intended to be used with `DecoupledIntQuant`."""
+
+    def __init__(
+            self,
+            stats_reduce_dim: int,
+            pre_zero_point_stats_input_view_shape_impl: Module,
+            pre_zero_point_shape: Optional[Tuple[int, ...]] = None) -> None:
+        super(PreZeroCenterZeroPoint, self).__init__()
+        self.stats_reduce_dim = stats_reduce_dim
+        self.stats_output_shape = pre_zero_point_shape
+        self.stats_input_view_shape_impl = pre_zero_point_stats_input_view_shape_impl
+
+    @brevitas.jit.script_method
+    def get_zero_center(self, x: Tensor) -> Tensor:
+        x = self.stats_input_view_shape_impl(x)
+        u = torch.mean(x, dim=self.stats_reduce_dim, keepdim=True)
+        z = -u.view(self.stats_output_shape)
+        return z
+
+    @brevitas.jit.script_method
+    def forward(self, x: Tensor, scale: Tensor, bit_width: Tensor) -> Tensor:
+        # NOTE: `DecoupledIntQuant` adds the `pre_zero_point` value to the scaled tensor,
+        # so this needs to return the negative of the scaled average value to perform
+        # pre-zero centering before rounding and clipping
+        z = self.get_zero_center(x) / scale  # need to scale the norm by s
+        return z

src/brevitas/function/ops.py

@@ -205,9 +205,10 @@ def max_float(exponent_bit_width: Tensor, mantissa_bit_width: Tensor, exponent_b
 
 def get_upper_bound_on_l1_norm(
         accumulator_bit_width: Tensor, input_bit_width: Tensor, input_is_signed: bool) -> Tensor:
-    """Calculate the upper bound on the l1-norm of the weights using the derivations from
-    `Quantized Neural Networks for Low-Precision Accumulation with Guaranteed Overflow Avoidance`
-    by I.Colbert, A.Pappalardo, and J.Petri-Koenig."""
+    """Calculate the upper bound on the l1-norm of the weights needed to guarantee overflow avoidance
+    for a given accumulator bit width and input representation using the derivations from
+    `A2Q: Accumulator-Aware Quantization with Guaranteed Overflow Avoidance` by I.Colbert,
+    A.Pappalardo, and J.Petri-Koenig. Note that this assumes integer quantization."""
     assert input_bit_width is not None, "A2Q relies on input bit-width."
     assert input_is_signed is not None, "A2Q relies on input sign."
     assert accumulator_bit_width is not None, "A2Q relies on accumulator bit-width."

src/brevitas/nn/utils.py

@@ -1,13 +1,10 @@
 # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
-from typing import Optional
-
 import torch
 from torch import Tensor
 from torch.nn import Parameter
 
-from brevitas.function.ops_ste import ceil_ste
 from brevitas.proxy.parameter_quant import BiasQuantProxyFromInjector
 from brevitas.proxy.parameter_quant import WeightQuantProxyFromInjector
 
@@ -79,51 +76,3 @@ def check_tensors_same_ptr(tensor_list):
         else:
             return False
     return all(p == pointers[0] for p in pointers)
-
-
-def calculate_min_accumulator_bit_width(
-        input_bit_width: Tensor,
-        input_is_signed: bool,
-        weight_max_l1_norm: Optional[Tensor] = None,
-        weight_bit_width: Optional[Tensor] = None,
-        n_elements: Optional[Tensor] = None,
-        min_val: Optional[float] = 1e-10):
-    """Using the closed-form bounds on accumulator bit-width as derived in `Quantized Neural Networks for Low-Precision Accumulation with Guaranteed Overflow
-    Avoidance` by I. Colbert, A. Pappalardo, and J. Petri-Koenig. This function returns the minimum accumulator bit-width that can be used without risk of
-    overflow. It supports both the data-type bound as well as the weight-level bound.
-
-    Args:
-        input_bit_width (Tensor): the bit-width of the inputs to the layer.
-        input_is_signed (bool): calculate statistics for normalizing weight parameter.
-        weight_max_l1_norm (Tensor): the maximum per-channel l1-norm of the weights.
-        weight_bit_width (Tensor): the bit-width of the weights to the layer.
-        n_elements (Tensor): the number of elements in the dot product.
-        min_val (float): the minimum value used for the l1-norm, used to avoid log2(0). Default: 1e-8.
-
-    Example (data-type bound):
-    >> acc_bit_width = calculate_min_accumulator_bit_width(input_bit_width, input_is_signed, weight_bit_width, n_elements)
-
-    Example (weight-level bound):
-    >> acc_bit_width = calculate_min_accumulator_bit_width(input_bit_width, input_is_signed, weight_max_l1_norm)
-    """
-    input_is_signed = float(input_is_signed)
-    # if the l1-norm of the weights is specified, then use the weight-level bound
-    if weight_max_l1_norm is not None:
-        assert isinstance(weight_max_l1_norm, (float, Tensor)), "The l1-norm of the weights needs to be a float or a torch.Tensor instance."
-        if isinstance(weight_max_l1_norm, Tensor):
-            assert weight_max_l1_norm.numel() == 1, "The minimum accumulator bit-width calculation currently only supports scalars."
-        weight_max_l1_norm = torch.clamp_min(weight_max_l1_norm, min_val)
-        input_is_signed = float(input_is_signed)
-        alpha = torch.log2(weight_max_l1_norm) + input_bit_width - input_is_signed
-    # else use the data-type bound
-    else:
-        assert isinstance(weight_bit_width, (float, Tensor)), "If weight_max_l1_norm is un-specified, weight_bit_width needs to be specified."
-        assert isinstance(n_elements, (float, Tensor)), "If weight_max_l1_norm is un-specified, n_elements needs to be specified."
-        if isinstance(n_elements, Tensor):
-            assert n_elements.numel() == 1, "The minimum accumulator bit-width calculation currently only supports scalars."
-        assert n_elements > 0, "There needs to be at least one element considered in this evaluation."
-        alpha = torch.log2(n_elements) + input_bit_width + weight_bit_width - input_is_signed - 1.
-    phi = lambda x: torch.log2(1. + pow(2., -x))
-    min_bit_width = alpha + phi(alpha) + 1.
-    min_bit_width = ceil_ste(min_bit_width)
-    return min_bit_width  # returns the minimum accumulator that can be used without risk of overflow

src/brevitas/quant/base.py

@@ -20,6 +20,7 @@
 from brevitas.core.restrict_val import FloatRestrictValue
 from brevitas.core.restrict_val import LogFloatRestrictValue
 from brevitas.core.scaling import AccumulatorAwareParameterPreScaling
+from brevitas.core.scaling import AccumulatorAwareZeroCenterParameterPreScaling
 from brevitas.core.scaling import IntScaling
 from brevitas.core.scaling import ParameterFromStatsFromParameterScaling
 from brevitas.core.scaling import ParameterPreScalingWeightNorm
@@ -38,6 +39,7 @@
 from brevitas.core.utils import SingleArgStatelessBuffer
 from brevitas.core.zero_point import ParameterFromRuntimeZeroPoint
 from brevitas.core.zero_point import ParameterFromStatsFromParameterZeroPoint
+from brevitas.core.zero_point import PreZeroCenterZeroPoint
 from brevitas.core.zero_point import StatsFromParameterZeroPoint
 from brevitas.core.zero_point import ZeroZeroPoint
 from brevitas.inject import ExtendedInjector
@@ -76,6 +78,7 @@
     'BatchQuantStatsScaling1d',
     'BatchQuantStatsScaling2d',
     'AccumulatorAwareWeightQuant',
+    'AccumulatorAwareZeroCenterWeightQuant',
     'MSESymmetricScale',
     'MSEAsymmetricScale',
     'MSEWeightZeroPoint',
@@ -400,6 +403,20 @@ def accumulator_bit_width_impl(accumulator_bit_width):
     float_to_int_impl = RoundToZeroSte  # required to ensure no upwards rounding violates constraints
 
 
+class AccumulatorAwareZeroCenterWeightQuant(AccumulatorAwareWeightQuant):
+    """Experimental zero-centered accumulator-aware weight quantized based on:
+    `A2Q+: Improving Accumulator-Aware Weight Quantization`.
+
+    When compared to A2Q, A2Q+ changes the following:
+    (1) added zero-centering constraint on the weights (i.e., `PreZeroCenterZeroPoint`)
+    (2) a more relaxed l1-norm bound that is derived in the referenced paper
+    """
+    pre_scaling_impl = AccumulatorAwareZeroCenterParameterPreScaling
+    pre_zero_point_impl = PreZeroCenterZeroPoint
+    pre_zero_point_shape = this.scaling_shape  # TODO: decouple zero_point from scaling
+    pre_zero_point_stats_input_view_shape_impl = this.scaling_stats_input_view_shape_impl
+
+
 class MSESubInjectorBase(ExtendedInjector):
 
     @value