Anatomy of a Quantizer#

What’s in a Quantizer?#

In a broad sense, a quantizer is anything that implements a quantization technique, and the flexibility of Brevitas means that there are different ways to do so.

However, to keep our terminology straight, we refer to a quantizer as a specific kind of way to implement quantization, the one preferred and adopted by default. That is, a quantizer is a subclass of a brevitas.inject.ExtendedInjector that carries a tensor_quant attribute, which points to an instance of a torch Module that implements quantization.

We have seen in previous tutorials quantizers being imported from brevitas.quant and passed on to quantized layers. We can easily very what we just said on one of them:

[1]:

from brevitas.inject import ExtendedInjector
from brevitas.quant.scaled_int import Int8ActPerTensorFloat

issubclass(Int8ActPerTensorFloat, ExtendedInjector)

[1]:

True

[2]:

Int8ActPerTensorFloat.tensor_quant

[2]:

RescalingIntQuant(
  (int_quant): IntQuant(
    (float_to_int_impl): RoundSte()
    (tensor_clamp_impl): TensorClamp()
    (delay_wrapper): DelayWrapper(
      (delay_impl): _NoDelay()
    )
    (input_view_impl): Identity()
  )
  (scaling_impl): ParameterFromRuntimeStatsScaling(
    (stats_input_view_shape_impl): OverTensorView()
    (stats): _Stats(
      (stats_impl): AbsPercentile()
    )
    (restrict_scaling): _RestrictValue(
      (restrict_value_impl): FloatRestrictValue()
    )
    (clamp_scaling): _ClampValue(
      (clamp_min_ste): ScalarClampMinSte()
    )
    (restrict_inplace_preprocess): Identity()
    (restrict_preprocess): Identity()
  )
  (int_scaling_impl): IntScaling()
  (zero_point_impl): ZeroZeroPoint(
    (zero_point): StatelessBuffer()
  )
  (msb_clamp_bit_width_impl): BitWidthConst(
    (bit_width): StatelessBuffer()
  )
)

Note how we said subclass and not instance. To understand why that’s the case, we have to understand what an ExtendedInjector is and why it’s used in the first place.

Quantization with auto-wiring Dependency Injection#

Pytorch has exploded in popularity thanks to its straightforward numpy-like define-by-run execution model. However, when it comes to applying quantization, this style of programming poses a problem.

Many quantization methods depend on making decisions based on the (in Pytorch terms) state_dict of the original floating-point model to finetune with quantization. However, when we instantiate a model in Pytorch we can’t know on the spot if a state_dict is going to be loaded a few lines of code later or not. Yet, because Pytorch is define-by-run, we need our model to work consistently both before and after a state_dict is possibly loaded. In a traditional scenario that wouldn’t pose a problem. However, with quantization in the loop, the way a quantizer is defined might change before and after a pretrained state_dict is loaded.

That means that we need a way to define our quantized model such that it can react appropriately in case the state_dict changes. In a Python-only world that wouldn’t be too hard. However, in order to mitigate the performance impact of quantization-aware training, Brevitas makes extended use of Pytorch’s JIT compiler for a custom subset of Python, TorchScript. That means that in most scenarios, when a state_dict is loaded, we need to recompile parts of the model. Because compilation in general is a lossy process, a TorchScript component cannot simply re-compile itself based on new input information.

We need then a way to declare a quantization method such that it can be re-initialized and JIT compiled any time the state_dict changes. Because we want to support arbitrarly-complex user-defined quantization algorithms, this method has to be generic, i.e. it cannot depend on the specifics of the quantization algorithm implemented.

Implementing a quantizer with an ExtendedInjector is a way to do so. Specifically, an ExtendedInjector extends an Injector from an older version (0.2.1) of the excellent dependency-injection library dependencies with support for a couple of extra features that are specific to Brevitas’ needs.

An Injector (and an ExtendedInjector) allows to take what might be a very complicated graph of interwined objects and turns it into a flat list of variables that are capable of auto-assembly by matching variable names to arguments names. This technique typically goes under the name of auto-wiring dependency injection.

In the context of Brevitas, the goal is gather all the modules and hyperparameters that contribute to a quantization implementation such that they can be re-assembled automatically on demand. What comes out of this process is a tensor_quant object.

A Practical Example: Binary Quantization#

To make things practical, let’s look at how we can implement a simple variant of binary quantization. All the components typically used to implement quantization can be found under brevitas.core. As mentioned before, Brevitas makes heavy use of TorchScript. In particular, all the components found under brevitas.core are implemented as ScriptModule that can be assembled together. The core ScriptModule that implements binarization can be found under brevitas.core.quant:

[3]:

import inspect
from IPython.display import Markdown, display

# helpers
def assert_with_message(condition):
    assert condition
    print(condition)

def pretty_print_source(source):
    display(Markdown('```python\n' + source + '\n```'))

[4]:

from brevitas.core.quant import BinaryQuant

source = inspect.getsource(BinaryQuant)
pretty_print_source(source)

class BinaryQuant(brevitas.jit.ScriptModule):
    """
    ScriptModule that implements scaled uniform binary quantization of an input tensor.
    Quantization is performed with :func:`~brevitas.function.ops_ste.binary_sign_ste`.

    Args:
        scaling_impl (Module): Module that returns a scale factor.
        quant_delay_steps (int): Number of training steps to delay quantization for. Default: 0

    Returns:
        Tuple[Tensor, Tensor, Tensor, Tensor]: Quantized output in de-quantized format, scale, zero-point, bit_width.

    Examples:
        >>> from brevitas.core.scaling import ConstScaling
        >>> binary_quant = BinaryQuant(ConstScaling(0.1))
        >>> inp = torch.Tensor([0.04, -0.6, 3.3])
        >>> out, scale, zero_point, bit_width = binary_quant(inp)
        >>> out
        tensor([ 0.1000, -0.1000,  0.1000])
        >>> scale
        tensor(0.1000)
        >>> zero_point
        tensor(0.)
        >>> bit_width
        tensor(1.)

    Note:
        Maps to quant_type == QuantType.BINARY == 'BINARY' == 'binary' when applied to weights in higher-level APIs.

    Note:
        Set env variable BREVITAS_JIT=1 to enable TorchScript compilation of this module.
    """

    def __init__(self, scaling_impl: Module, signed: bool = True, quant_delay_steps: int = 0):
        super(BinaryQuant, self).__init__()
        assert signed, "Unsigned binary quant not supported"
        self.scaling_impl = scaling_impl
        self.bit_width = BitWidthConst(1)
        self.zero_point = StatelessBuffer(torch.tensor(0.0))
        self.delay_wrapper = DelayWrapper(quant_delay_steps)

    @brevitas.jit.script_method
    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        scale = self.scaling_impl(x)
        y = binary_sign_ste(x) * scale
        y = self.delay_wrapper(x, y)
        return y, scale, self.zero_point(), self.bit_width()

The implementation is quite simple. Apart from quant_delay_steps, which allows to delay quantization by a certain number of training steps (default = 0), the only other argument that BinaryQuant accepts is an implementation to compute the scale factor. bit_width is fixed to 1 and zero-point is fixed to 0.

We pick as scale factor implementation a ScriptModule called ParameterScaling, which implements a learned parameter with user-defined initialization. It can be found under brevitas.core.scaling:

[5]:

from brevitas.core.scaling import ParameterScaling

Manual Binary Quantization#

As a first step, we simply instantiate BinaryQuant with ParameterScaling using scaling_init equal 0.1 and we call it on a random floating-point input tensor:

[6]:

import torch

# set seed for notebook
torch.manual_seed(0)

manual_tensor_quant = BinaryQuant(scaling_impl=ParameterScaling(scaling_init=0.1))
manual_tensor_quant(torch.randn(4, 4))

[6]:

(tensor([[-0.1000, -0.1000, -0.1000, -0.1000],
         [ 0.1000,  0.1000, -0.1000, -0.1000],
         [ 0.1000, -0.1000,  0.1000,  0.1000],
         [ 0.1000,  0.1000,  0.1000, -0.1000]], grad_fn=<MulBackward0>),
 tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>),
 tensor(0.),
 tensor(1.))

Nothing too surprising here, as expected the tensor is binarized with the scale factor we defined. Note however how manual_tensor_quant is returning a tuple and not a QuantTensor. This is because support for custom data structures in TorchScript is still quite limited, so QuantTensor are allocated only in Python-world abstractions.

Binary Quantization with an ExtendedInjector#

Let’s now declare tensor_quant through an ExtendedInjector:

[7]:

from brevitas.inject import ExtendedInjector

class MyBinaryQuantizer(ExtendedInjector):
    tensor_quant = BinaryQuant
    scaling_impl=ParameterScaling
    scaling_init=0.1

inj_tensor_quant = MyBinaryQuantizer.tensor_quant
inj_tensor_quant(torch.randn(4, 4))

[7]:

(tensor([[-0.1000, -0.1000,  0.1000,  0.1000],
         [ 0.1000, -0.1000, -0.1000,  0.1000],
         [ 0.1000, -0.1000, -0.1000,  0.1000],
         [ 0.1000,  0.1000,  0.1000, -0.1000]], grad_fn=<MulBackward0>),
 tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>),
 tensor(0.),
 tensor(1.))

Any time MyBinaryQuantizer.tensor_quant is called, a new instance of BinaryQuant is created. Note how the attributes of MyBinaryQuantizer are designed to match the name of the arguments of each other, except for tensor_quant, which is what we are interested in retrieving from the outside.

Inheritance and Composition of Quantizers#

The advantage of expressing a quantizer through a Python class also means that we can leverage both inheritance and composition. So for example we can inherit from MyBinaryQuantizer and override scaling_init with a new value:

[8]:

class MyChildBinaryQuantizer(MyBinaryQuantizer):
    scaling_init=1.0

child_inj_tensor_quant = MyChildBinaryQuantizer.tensor_quant
child_inj_tensor_quant(torch.randn(4, 4))

[8]:

(tensor([[-1.,  1., -1.,  1.],
         [ 1.,  1.,  1.,  1.],
         [-1.,  1., -1.,  1.],
         [ 1.,  1., -1., -1.]], grad_fn=<MulBackward0>),
 tensor(1., grad_fn=<AbsBinarySignGradFnBackward>),
 tensor(0.),
 tensor(1.))

Or we can leverage composition by assembling together various classes containing different pieces of a quantizer:

[9]:

class MyBinaryImpl(ExtendedInjector):
    tensor_quant = BinaryQuant

class MyScalingImpl(ExtendedInjector):
    scaling_impl=ParameterScaling
    scaling_init=0.1

class MyComposedBinaryQuantizer(MyBinaryImpl, MyScalingImpl):
    pass

comp_inj_tensor_quant = MyComposedBinaryQuantizer.tensor_quant
comp_inj_tensor_quant(torch.randn(4, 4))

[9]:

(tensor([[-0.1000,  0.1000,  0.1000,  0.1000],
         [-0.1000, -0.1000,  0.1000, -0.1000],
         [-0.1000,  0.1000,  0.1000,  0.1000],
         [-0.1000, -0.1000,  0.1000, -0.1000]], grad_fn=<MulBackward0>),
 tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>),
 tensor(0.),
 tensor(1.))

Interfacing a Quantizer with a Quantized Layer#

Before we can pass the quantizer to a quantized layer such as QuantConv2d, we need a last component to define, a proxy. A proxy (found under brevitas.proxy) is an nn.Module that serve as interface between a quantizer and a quantized layer.

While a quantizer lives mostly in JIT-land, a proxy lives mostly in Python-land, and as such can afford much more flexibility. Proxies take care of returning a QuantTensor and re-initializing the output of quantizer whenever a new state_dict is loaded.

Proxies are specific to the kind of tensor being quantized, as in weights vs biases vs activations. For convenience, they are declared as part of the quantizer itself under the attribute proxy_class. For example, for weights we can use WeightQuantProxyFromInjector:

[10]:

from brevitas.proxy import WeightQuantProxyFromInjector

class MyBinaryWeightQuantizer(MyBinaryQuantizer):
    proxy_class = WeightQuantProxyFromInjector

We can now use MyBinaryWeightQuantizer as the weight quantizer of a layer:

[11]:

from brevitas.nn import QuantConv2d

binary_weight_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MyBinaryWeightQuantizer)
try:
    quant_weight = binary_weight_quant_conv.quant_weight()
except TypeError:
    pass

Note however that we cannot compute the quantized weight, as the signed attribute is None.

signed is one of those attributes that in the case of binary quantization has to be explicitly defined by the user. The idea is that it informs the proxy on whether the value generated by our quantizer should be considered signed or not. We can do so by simply setting it in the quantizer:

[12]:

class MySignedBinaryWeightQuantizer(MyBinaryWeightQuantizer):
    signed = True

binary_weight_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MySignedBinaryWeightQuantizer)
signed_quant_weight = binary_weight_quant_conv.quant_weight()
signed_quant_weight

[12]:

IntQuantTensor(value=tensor([[[[ 0.1000,  0.1000, -0.1000],
          [ 0.1000, -0.1000, -0.1000],
          [-0.1000, -0.1000, -0.1000]],

         [[-0.1000, -0.1000,  0.1000],
          [-0.1000, -0.1000,  0.1000],
          [-0.1000,  0.1000, -0.1000]],

         [[ 0.1000,  0.1000,  0.1000],
          [ 0.1000,  0.1000, -0.1000],
          [-0.1000, -0.1000,  0.1000]]],


        [[[-0.1000, -0.1000, -0.1000],
          [-0.1000, -0.1000, -0.1000],
          [ 0.1000,  0.1000, -0.1000]],

         [[ 0.1000, -0.1000, -0.1000],
          [-0.1000, -0.1000, -0.1000],
          [-0.1000, -0.1000, -0.1000]],

         [[ 0.1000, -0.1000,  0.1000],
          [-0.1000, -0.1000,  0.1000],
          [-0.1000,  0.1000,  0.1000]]]], grad_fn=<MulBackward0>), scale=tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

[13]:

assert_with_message(signed_quant_weight.is_valid)

True

And now the quant weights are valid.

When we want to add or override an single attribute of a quantizer passed to a layer, defining a whole new quantizer can be too verbose. There is a simpler syntax to achieve the same goal. Let’s say we want to have add the signed attribute to MyBinaryQuantizer, as we just did. We could have also simply done the following:

[14]:

small_scale_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MyBinaryWeightQuantizer, weight_signed=True)
small_scale_quant_conv.quant_weight()

[14]:

IntQuantTensor(value=tensor([[[[-0.1000, -0.1000,  0.1000],
          [-0.1000,  0.1000, -0.1000],
          [ 0.1000,  0.1000,  0.1000]],

         [[ 0.1000, -0.1000, -0.1000],
          [-0.1000,  0.1000,  0.1000],
          [ 0.1000, -0.1000,  0.1000]],

         [[-0.1000, -0.1000, -0.1000],
          [ 0.1000, -0.1000, -0.1000],
          [ 0.1000,  0.1000, -0.1000]]],


        [[[-0.1000, -0.1000,  0.1000],
          [-0.1000,  0.1000,  0.1000],
          [ 0.1000, -0.1000, -0.1000]],

         [[ 0.1000, -0.1000, -0.1000],
          [ 0.1000, -0.1000, -0.1000],
          [ 0.1000, -0.1000, -0.1000]],

         [[ 0.1000,  0.1000,  0.1000],
          [-0.1000,  0.1000, -0.1000],
          [-0.1000, -0.1000, -0.1000]]]], grad_fn=<MulBackward0>), scale=tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

What we did was to take the name of the attribute signed, add the prefix weight_, and pass it as a keyword argument to QuantConv2d. What happens in the background is that the keyword arguments prefixed with weight_ are set as attributes of weight_quant, possibly overriding any pre-existing value. The same principle applies to input_, output_ and bias_.

This is the reason why, as it was mentioned in the first tutorial, quantized layers can accept arbitrary keyword arguments. It’s really just a way to support different styles of syntax when defining a quantizer.

Passing a custom quantizer to QuantIdentity#

We can do a similar thing with quantized activations:

[15]:

from brevitas.proxy import ActQuantProxyFromInjector
from brevitas.nn import QuantIdentity

class MySignedBinaryActQuantizer(MyBinaryQuantizer):
    proxy_class = ActQuantProxyFromInjector
    signed = True

binary_relu = QuantIdentity(act_quant=MySignedBinaryActQuantizer, return_quant_tensor=True)
binary_relu(torch.randn(4, 4))

/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/_tensor.py:1255: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /opt/conda/conda-bld/pytorch_1670525541990/work/c10/core/TensorImpl.h:1758.)
  return super(Tensor, self).rename(names)

[15]:

IntQuantTensor(value=tensor([[-0.1000,  0.1000, -0.1000, -0.1000],
        [ 0.1000,  0.1000,  0.1000,  0.1000],
        [-0.1000, -0.1000,  0.1000, -0.1000],
        [-0.1000,  0.1000, -0.1000,  0.1000]], grad_fn=<MulBackward0>), scale=tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

So there isn’t really much difference between a quantizer for weights and a quantizer for activations, they are just wrapped by different proxies. Also, with activations a prefix is not required when passing keyword arguments. For example, when can override the existing scaling_init defined in MyBinaryQuantizer with a new value passed in as a keywork argument:

[16]:

small_scale_binary_identity = QuantIdentity(
    act_quant=MySignedBinaryActQuantizer, scaling_init=0.001, return_quant_tensor=True)
small_scale_binary_identity(torch.randn(4, 4))

[16]:

IntQuantTensor(value=tensor([[ 0.0010, -0.0010, -0.0010,  0.0010],
        [ 0.0010,  0.0010,  0.0010,  0.0010],
        [ 0.0010, -0.0010,  0.0010,  0.0010],
        [ 0.0010, -0.0010, -0.0010, -0.0010]], grad_fn=<MulBackward0>), scale=tensor(0.0010, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

A Custom Quantizer initialized with Weight Statistics#

So far we have seen use-cases where an ExtendedInjector provides, at best, a different kind of syntax to define a quantizer, without any particular other advantage. Let’s now make things a bit more complicated to show the sort of situations where it really shines.

Let’s say we want to define a binary weight quantizer where scaling_impl is still ParameterScaling. However, instead of being user-defined, we want scaling_init to be the maximum value found in the weight tensor of the quantized layer. To support this sort of use cases where the quantizer depends on the layer, a quantized layer automatically passes itself to all its quantizers under the name of module. With only a few lines of code then, we can achieve our goal:

[17]:

from brevitas.inject import value

class ParamFromMaxWeightQuantizer(MySignedBinaryWeightQuantizer):

    @value
    def scaling_init(module):
        return module.weight.abs().max()

Note how we are leveraging the @value decorator to define a function that is executed at dependency-injection (DI) time. This kind of behaviour is similar in spirit to defining a @property instead of an attribute, with the difference that a @value function can depend on other attributes of the Injector, which are automatically passed in as arguments of the function during DI.

Let’s now pass the quantizer to a QuantConv2d and retrieve its quantized weights:

[18]:

param_from_max_quant_conv = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)
param_from_max_quant_conv.quant_weight()

[18]:

IntQuantTensor(value=tensor([[[[ 0.1820, -0.1820, -0.1820],
          [ 0.1820,  0.1820,  0.1820],
          [-0.1820, -0.1820, -0.1820]],

         [[ 0.1820, -0.1820, -0.1820],
          [ 0.1820, -0.1820, -0.1820],
          [ 0.1820,  0.1820, -0.1820]],

         [[-0.1820,  0.1820,  0.1820],
          [ 0.1820, -0.1820,  0.1820],
          [-0.1820, -0.1820, -0.1820]]],


        [[[ 0.1820,  0.1820, -0.1820],
          [-0.1820, -0.1820,  0.1820],
          [-0.1820,  0.1820, -0.1820]],

         [[-0.1820,  0.1820, -0.1820],
          [ 0.1820,  0.1820, -0.1820],
          [ 0.1820,  0.1820,  0.1820]],

         [[-0.1820, -0.1820, -0.1820],
          [ 0.1820, -0.1820,  0.1820],
          [ 0.1820, -0.1820,  0.1820]]]], grad_fn=<MulBackward0>), scale=tensor(0.1820, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

Indeed we can verify that weight_quant.scale() is equal to weight.abs().max():

[19]:

assert_with_message((param_from_max_quant_conv.weight_quant.scale() == param_from_max_quant_conv.weight.abs().max()).item())

True

Let’s say now that we want to load a pretrained floating-point weight tensor on top of our quantized model. We simuate this scenario by defining a separate nn.Conv2d layer with the same weight shape:

[20]:

from torch import nn

float_conv = nn.Conv2d(3, 2, (3, 3))
float_conv.weight.abs().max()

[20]:

tensor(0.1924, grad_fn=<MaxBackward1>)

and then we load it on top of param_from_max_quant_conv:

[21]:

param_from_max_quant_conv.load_state_dict(float_conv.state_dict())

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[21], line 1
----> 1 param_from_max_quant_conv.load_state_dict(float_conv.state_dict())

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py:1671, in Module.load_state_dict(self, state_dict, strict)
   1666         error_msgs.insert(
   1667             0, 'Missing key(s) in state_dict: {}. '.format(
   1668                 ', '.join('"{}"'.format(k) for k in missing_keys)))
   1670 if len(error_msgs) > 0:
-> 1671     raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
   1672                        self.__class__.__name__, "\n\t".join(error_msgs)))
   1673 return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for QuantConv2d:
        Missing key(s) in state_dict: "weight_quant.tensor_quant.scaling_impl.value".

Ouch, we get an error. This is because ParameterScaling contains a learned torch.nn.Parameter, and Pytorch expects all learned parameters of a model to be contained in a state_dict that is being loaded. We can work around the issue by either setting the IGNORE_MISSING_KEYS config flag in Brevitas, or by passing strict=False to load_state_dict. We go with the former as setting strict=False is too forgiving to other kind of problems:

[22]:

from brevitas import config
config.IGNORE_MISSING_KEYS = True

param_from_max_quant_conv.load_state_dict(float_conv.state_dict())

[22]:

<All keys matched successfully>

Note that we could have also achieve the same goal by setting the env variable BREVITAS_IGNORE_MISSING_KEYS=1.

And now if we take a look at the quantized weights again:

[23]:

param_from_max_quant_conv.quant_weight()

[23]:

IntQuantTensor(value=tensor([[[[ 0.1924,  0.1924, -0.1924],
          [ 0.1924,  0.1924,  0.1924],
          [ 0.1924,  0.1924,  0.1924]],

         [[-0.1924, -0.1924, -0.1924],
          [ 0.1924,  0.1924, -0.1924],
          [ 0.1924,  0.1924,  0.1924]],

         [[ 0.1924,  0.1924, -0.1924],
          [-0.1924,  0.1924,  0.1924],
          [-0.1924,  0.1924,  0.1924]]],


        [[[ 0.1924, -0.1924,  0.1924],
          [-0.1924,  0.1924, -0.1924],
          [ 0.1924,  0.1924,  0.1924]],

         [[-0.1924,  0.1924, -0.1924],
          [ 0.1924, -0.1924, -0.1924],
          [-0.1924,  0.1924,  0.1924]],

         [[ 0.1924,  0.1924, -0.1924],
          [-0.1924, -0.1924, -0.1924],
          [ 0.1924, -0.1924, -0.1924]]]], grad_fn=<MulBackward0>), scale=tensor(0.1924, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

We see that, as expected, the scale factor has been updated to the new weight.abs().max().

What happens internally is that after load_state_dict is called on the layer, ParamFromMaxWeightQuantizer.tensor_quant gets called again to re-initialize BinaryQuant, and in turn ParameterScaling is re-initialized with a new scaling_init value computed based on the updated module.weight tensor. This whole process wouldn’t have been possible without an ExtendedInjector behind it.

Dealing with Weight Initialization#

There is a type of situation that Brevitas cannot deal with automatically. That is, when the initialization of the quantizer depends on the layer to which it is applied (like with the ParamFromMaxWeightQuantizer or SharedParamFromMeanWeightQuantizer quantizers), but the layer gets modified after it is initialized.

The typical example is with weight initialization when training from scratch (so rather than loading from a floating-point state_dict):

[29]:

quant_conv_w_init = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)
torch.nn.init.uniform_(quant_conv_w_init.weight)

assert_with_message(not (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())

True

We can see how the scale factor is not initialized correctly anymore. In this case we can simply trigger re-initialization of the weight quantizer manually:

[30]:

quant_conv_w_init.weight_quant.init_tensor_quant()

assert_with_message((quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())

True

Note: because the way weights are initialized is often the same as to how an optimizer performs the weight update step, there are currently no plans to try to perform re-initialization automatically (as it happens e.g. when a state_dict is loaded) since it wouldn’t be possible to distinguish between the two scenarios.

Building a Custom Quantization API#

Finally, let’s go through an even more complicated example. We are going to look at a scenario that illustrates the differences between a standard Injector (implemented in the dependencies library) and our ExtendedInjector extension.

Let’s say we want to build two quantizers for respectively weights and activations and build a simple API on top of them. In particular, we want to be able to switch between BinaryQuant and ClampedBinaryQuant (a variant of binary quantization with clamping), and we want to optionally perform per-channel scaling. To do so, we are going to implement the controlling logic through a hierarchy of ExtendedInjector, leaving two boolean flags exposed as arguments of the quantizers, with the idea then that the flags can be set through keyword arguments of the respective quantized layers.

We can go as follows:

[31]:

from brevitas.core.quant import ClampedBinaryQuant
from brevitas.proxy import WeightQuantProxyFromInjector, ActQuantProxyFromInjector
from brevitas.inject import this


class CommonQuantizer(ExtendedInjector):
    scaling_impl = ParameterScaling
    signed=True

    @value
    def tensor_quant(is_clamped):
        # returning a class to auto-wire from a value function
        # wouldn't be allowed in a standard Injector
        if is_clamped:
            return ClampedBinaryQuant
        else:
            return BinaryQuant

    @value
    def scaling_shape(scaling_per_output_channel):
        if scaling_per_output_channel:
            # returning this.something from a value function
            # wouldn't be allowed in a standard Injector
            return this.per_channel_broadcastable_shape
        else:
            return ()


class AdvancedWeightQuantizer(CommonQuantizer):
    proxy_class = WeightQuantProxyFromInjector

    @value
    def per_channel_broadcastable_shape(module):
        return (module.weight.shape[0], 1, 1, 1)

    @value
    def scaling_init(module, scaling_per_output_channel):
        if scaling_per_output_channel:
            num_ch = module.weight.shape[0]
            return module.weight.abs().view(num_ch, -1).max(dim=1)[0].view(-1, 1, 1, 1)
        else:
            return module.weight.abs().max()


class AdvancedActQuantizer(CommonQuantizer):
    scaling_init = 0.01
    proxy_class = ActQuantProxyFromInjector

There are a bunch of things going on here to unpack.

The first one is that a @value function can return a class to auto-wire and inject, as seen in the definition of tensor_quant. This wouldn’t normally be possible with a standard Injector, but it’s possible with an ExtendedInjector. This way we can switch between different implementations of tensor_quant.

The second one is the special object this. this is already present in the dependencies library, and it’s used as a way to retrieve attributes of the quantizer from within the quantizer itself. However, normally it wouldn’t be possible to return a reference to this from a @value function. Again this is something that only a ExtendedInjector supports, and it allows to chain different attributes in a way such that the chained values are computed only when necessary.

Let’s see the quantizers applied to a layer:

[32]:

per_channel_quant_conv = QuantConv2d(
    3, 2, (3, 3),
    weight_quant=AdvancedWeightQuantizer,
    weight_is_clamped=False,
    weight_scaling_per_output_channel=True)
per_channel_quant_conv.quant_weight()

[32]:

IntQuantTensor(value=tensor([[[[ 0.1612, -0.1612, -0.1612],
          [-0.1612, -0.1612, -0.1612],
          [ 0.1612,  0.1612,  0.1612]],

         [[-0.1612,  0.1612, -0.1612],
          [-0.1612,  0.1612,  0.1612],
          [-0.1612, -0.1612,  0.1612]],

         [[-0.1612,  0.1612,  0.1612],
          [ 0.1612,  0.1612, -0.1612],
          [ 0.1612,  0.1612,  0.1612]]],


        [[[ 0.1924,  0.1924,  0.1924],
          [-0.1924, -0.1924,  0.1924],
          [-0.1924,  0.1924, -0.1924]],

         [[ 0.1924, -0.1924,  0.1924],
          [ 0.1924,  0.1924, -0.1924],
          [ 0.1924, -0.1924, -0.1924]],

         [[-0.1924, -0.1924,  0.1924],
          [ 0.1924, -0.1924, -0.1924],
          [ 0.1924, -0.1924,  0.1924]]]], grad_fn=<MulBackward0>), scale=tensor([[[[0.1612]]],


        [[[0.1924]]]], grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

As expected the weight scale is now a vector. Everything we said so far about quantizers still applies, so for example we can load the floating-point state dict we defined before and observe how it triggers an update of the weight scale:

[33]:

per_channel_quant_conv.load_state_dict(float_conv.state_dict())
per_channel_quant_conv.quant_weight()

[33]:

IntQuantTensor(value=tensor([[[[ 0.1924,  0.1924, -0.1924],
          [ 0.1924,  0.1924,  0.1924],
          [ 0.1924,  0.1924,  0.1924]],

         [[-0.1924, -0.1924, -0.1924],
          [ 0.1924,  0.1924, -0.1924],
          [ 0.1924,  0.1924,  0.1924]],

         [[ 0.1924,  0.1924, -0.1924],
          [-0.1924,  0.1924,  0.1924],
          [-0.1924,  0.1924,  0.1924]]],


        [[[ 0.1899, -0.1899,  0.1899],
          [-0.1899,  0.1899, -0.1899],
          [ 0.1899,  0.1899,  0.1899]],

         [[-0.1899,  0.1899, -0.1899],
          [ 0.1899, -0.1899, -0.1899],
          [-0.1899,  0.1899,  0.1899]],

         [[ 0.1899,  0.1899, -0.1899],
          [-0.1899, -0.1899, -0.1899],
          [ 0.1899, -0.1899, -0.1899]]]], grad_fn=<MulBackward0>), scale=tensor([[[[0.1924]]],


        [[[0.1899]]]], grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

In this case we have a per-channel quantizer, so the original floating-point weight tensor is now quantized per channel.

Similarly, we can apply our custom activation quantizer to e.g. a QuantIdentity layer:

[34]:

from brevitas.nn import QuantIdentity

quant_identity = QuantIdentity(
    act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=False)
quant_identity(torch.randn(4, 4))

[34]:

tensor([[ 0.0100,  0.0100, -0.0100,  0.0100],
        [ 0.0100,  0.0100, -0.0100,  0.0100],
        [-0.0100,  0.0100, -0.0100, -0.0100],
        [ 0.0100, -0.0100, -0.0100, -0.0100]], grad_fn=<MulBackward0>)

Note how AdvancedActQuantizer doesn’t define a per_channel_broadcastable_shape, yet no errors are triggered. This is because this.per_channel_broadcastable_shape is required only when scaling_per_output_channel is True, while in this case scaling_per_output_channel is False. Let’ try to set it to True then:

[35]:

from brevitas.nn import QuantIdentity

quant_identity = QuantIdentity(
    act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True)

---------------------------------------------------------------------------
DependencyError                           Traceback (most recent call last)
Cell In[35], line 3
      1 from brevitas.nn import QuantIdentity
----> 3 quant_identity = QuantIdentity(
      4     act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True)

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_activation.py:113, in QuantIdentity.__init__(self, act_quant, return_quant_tensor, **kwargs)
    108 def __init__(
    109         self,
    110         act_quant: Optional[ActQuantType] = Int8ActPerTensorFloat,
    111         return_quant_tensor: bool = False,
    112         **kwargs):
--> 113     QuantNLAL.__init__(
    114         self,
    115         input_quant=None,
    116         act_impl=None,
    117         passthrough_act=True,
    118         act_quant=act_quant,
    119         return_quant_tensor=return_quant_tensor,
    120         **kwargs)

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_layer.py:34, in QuantNonLinearActLayer.__init__(self, act_impl, passthrough_act, input_quant, act_quant, return_quant_tensor, **kwargs)
     32 QuantLayerMixin.__init__(self, return_quant_tensor)
     33 QuantInputMixin.__init__(self, input_quant, **kwargs)
---> 34 QuantNonLinearActMixin.__init__(self, act_impl, passthrough_act, act_quant, **kwargs)

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/act.py:66, in QuantNonLinearActMixin.__init__(self, act_impl, passthrough_act, act_quant, act_proxy_prefix, act_kwargs_prefix, **kwargs)
     55 def __init__(
     56         self,
     57         act_impl: Optional[Type[Module]],
   (...)
     61         act_kwargs_prefix='',
     62         **kwargs):
     63     prefixed_kwargs = {
     64         act_kwargs_prefix + 'act_impl': act_impl,
     65         act_kwargs_prefix + 'passthrough_act': passthrough_act}
---> 66     QuantProxyMixin.__init__(
     67         self,
     68         quant=act_quant,
     69         proxy_prefix=act_proxy_prefix,
     70         kwargs_prefix=act_kwargs_prefix,
     71         proxy_protocol=ActQuantProxyProtocol,
     72         none_quant_injector=NoneActQuant,
     73         **prefixed_kwargs,
     74         **kwargs)

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/base.py:48, in QuantProxyMixin.__init__(self, quant, proxy_protocol, none_quant_injector, proxy_prefix, kwargs_prefix, **kwargs)
     46     quant_injector = quant
     47     quant_injector = quant_injector.let(**filter_kwargs(kwargs_prefix, kwargs))
---> 48     quant = quant_injector.proxy_class(self, quant_injector)
     49 else:
     50     if not isinstance(quant, proxy_protocol):

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/runtime_quant.py:198, in ActQuantProxyFromInjector.__init__(self, quant_layer, quant_injector)
    197 def __init__(self, quant_layer, quant_injector):
--> 198     super().__init__(quant_layer, quant_injector)
    199     self.cache_class = _CachedIO

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/runtime_quant.py:93, in ActQuantProxyFromInjectorBase.__init__(self, quant_layer, quant_injector)
     92 def __init__(self, quant_layer, quant_injector):
---> 93     QuantProxyFromInjector.__init__(self, quant_layer, quant_injector)
     94     ActQuantProxyProtocol.__init__(self)
     95     self.is_passthrough_act = _is_passthrough_act(quant_injector)

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/quant_proxy.py:80, in QuantProxyFromInjector.__init__(self, quant_layer, quant_injector)
     78 # Use a normal list and not a ModuleList since this is a pointer to parent modules
     79 self.tracked_module_list = []
---> 80 self.add_tracked_module(quant_layer)
     81 self.disable_quant = False
     82 # Torch.compile compatibility requires this

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/quant_proxy.py:120, in QuantProxyFromInjector.add_tracked_module(self, module)
    118     self.tracked_module_list.append(module)
    119     self.update_tracked_modules()
--> 120     self.init_tensor_quant()
    121 else:
    122     raise RuntimeError("Trying to add None as a parent module.")

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/runtime_quant.py:127, in ActQuantProxyFromInjectorBase.init_tensor_quant(self)
    126 def init_tensor_quant(self):
--> 127     tensor_quant = self.quant_injector.tensor_quant
    128     if 'act_impl' in self.quant_injector:
    129         act_impl = self.quant_injector.act_impl

    [... skipping hidden 1 frame]

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/_dependencies/this.py:51, in _ThisSpec.__call__(self, __self__)
     49 if kind == ".":
     50     try:
---> 51         result = getattr(result, symbol)
     52     except DependencyError:
     53         message = (
     54             "You tried to shift this more times than Injector has levels"
     55         )

File /proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/inject/__init__.py:129, in _ExtendedInjectorType.__getattr__(cls, attrname)
    126     else:
    127         message = "{!r} can not resolve attribute {!r}".format(
    128             cls.__name__, current_attr)
--> 129     raise DependencyError(message)
    131 marker, attribute, args, have_defaults = spec
    133 if set(args).issubset(cached):

DependencyError: 'AdvancedActQuantizer' can not resolve attribute 'per_channel_broadcastable_shape'

As expected we get an error saying that the quantizer cannot resolve per_channel_broadcastable_shape. If we pass it in then we can get a per-channel quantizer:

[36]:

quant_identity = QuantIdentity(
    act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True,
    per_channel_broadcastable_shape=(4, 1), return_quant_tensor=True)
quant_identity(torch.randn(4, 4))

[36]:

IntQuantTensor(value=tensor([[ 0.0100,  0.0100,  0.0100, -0.0100],
        [ 0.0100, -0.0100, -0.0100,  0.0100],
        [-0.0100, -0.0100, -0.0100,  0.0100],
        [ 0.0100,  0.0100,  0.0100,  0.0100]], grad_fn=<MulBackward0>), scale=tensor([[0.0100],
        [0.0100],
        [0.0100],
        [0.0100]], grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))

We have seen how powerful dependency injection is. In a way, it’s even too expressive. For users that are not interesting in building completely custom quantizers, it can be hard to make sense of how the various components available under brevitas.core can be assembled together according to best practices.