Minifloat and Groupwise quantization#

This notebook shows some practical use cases for minifloat and groupwise quantization.

Brevitas supports a wide combination of float quantization, including the OCP and FNUZ FP8 standard. It is possible to define any combination of exponent/mantissa bitwidth, as well as exponent bias.

Similarly, MX quantization is supported as general groupwise quantization on top of integer/minifloat datatype. This allows to any general groupwise quantization, including MXInt and MXFloat standards.

This tutorial shows how to instantiate and use some of the most interesting quantizers for minifloat and groupwise quantization

Minifloat (FP8 and lower)#

Brevitas offers some pre-defined quantizers for minifloat quantization, including OCP and FNUZ standards, which can be further customized according to the specific use case. The general naming structure for the quantizers is the following:

Fp\<Bitwidth\>\<Standard\>Weight\<Scaling\>Float

Where Bitwidth can be either empty or e4m3/e5m2, Standard can be empty or OCP/FNUZ, Scaling can be empty or PerTensor/PerChannel.

If Bitwidth is empty, the user must set it with kwargs or by subclassing the quantizers. Once the bitwidth is defined, the correct values for inf/nan are automatically defined based on the Standard. If a non-valid OCP bitwidth is set (e.g., e6m1), then no inf/nan values will be selected and the corresponding quantizer is not standard-compliant.

Standard allows to pick among the two main FP8 standard options; moreover, if not specified, Brevitas offers the possibility of doing minifloat quantization without necessarily reserving values for inf/nan representation. This allows to use the maximum available range, since often in quantization, values that exceed the quantization range saturate to maximum rather than going to inf/nan. FNUZ quantizers need to have saturating=True.

The Scaling options defines whether the quantization is scaled or unscaled. In the unscaled case, the scale factor for quantization is fixed to one, otherwise it can be set using any of the methods that Brevitas includes (e.g., statistics, learned, etc.)

Please keep in mind that not all combinations of the above options might be pre-defined and this serves mostly as indications of what Brevitas supports. It is possible, following the same structure of the available quantizers, to define new ones that fit your needs.

Similar considerations can be extended for activation quantization.

[ ]:

from brevitas.quant.float_base import Fp8e4m3Mixin
from brevitas.quant.mx_quant_ocp import MXFloat8e4m3Weight
from brevitas.quant.float_quant_ocp import FpOCPWeightPerTensorFloat, FpOCPActPerTensorFloat
from brevitas.quant.mx_quant_ocp import MXFloat8e4m3Act
import brevitas.nn as qnn
import torch.nn as nn
import torch
from brevitas.quant_tensor import FloatQuantTensor

class OCPFP8Weight(FpOCPWeightPerTensorFloat, Fp8e4m3Mixin):
    pass


class OCPFP8Act(FpOCPActPerTensorFloat, Fp8e4m3Mixin):
    pass


class FP8Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = qnn.QuantConv2d(32, 64, 3, weight_quant=OCPFP8Weight, input_quant=OCPFP8Act)

    def forward(self, x):
        return self.conv(x)

ocp_fp8_model = FP8Model()
x = torch.randn(1, 32, 8, 8)
ocp_fp8_model.eval()
o = ocp_fp8_model(x)

intermediate_input = ocp_fp8_model.conv.input_quant(x)
assert isinstance(intermediate_input, FloatQuantTensor)
assert isinstance(ocp_fp8_model.conv.quant_weight(), FloatQuantTensor)

/home/giuseppe/Documents/git/working_brevitas/.venv/lib/python3.11/site-packages/torch/_tensor.py:1644: 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 /pytorch/c10/core/TensorImpl.h:1935.)
  return super().rename(names)

Groupwise quantization (MXInt/MXFloat)#

Groupwise quantization is built on top of integer/minifloat quantization, with special considerations to accomodate for the groupwise scaling.

Compared to Int/Float QuantTensor, the main difference of their groupwise equivalent is that value, scale, and zero_point are not direct attributes anymore but properties. The new attributes are value_, scale_, and zero_point_.

The reason for this is shaping. When quantizing a tensor with shapes [O, I], where O is output channel and I is input channel, with groupsize k, groupwise quantization is normally represented as follow:

  • Tensor with shapes [O, I/k, k]

  • Scales with shapes [O, I/k, 1]

  • Zero point same as scale

The alternative to this representation is to have all three tensors with shapes [O,I], with a massive increase in memory utilization, especially with QAT + gradients.

The underscored attributes will have the compressed shapes, while the properties (non-underscored naming) will dynamically compute the expanded version of the property. This means:

quant_tensor.scale_.shape
# This will print [O, I/k, 1]
quant_tensor.scale.shape
# This will print [O, I]

With respect to pre-defined quantizers, Brevitas offers several Groupwise and MX options. The main difference between the two is that MX is restricted to group_size=32 and the scale factor must be a power-of-2. The user can override these settings but the corresponding output won’t be MX compliant.

Another difference is that MXFloat relies on the OCP format as underlying data type, while generic groupwise float relies on the non-standard minifloat representation explained above.

Finally, the general groupwise scaling relies on float scales.

[2]:

from brevitas.quant_tensor import GroupwiseFloatQuantTensor


class MXFloat8Weight(MXFloat8e4m3Weight, Fp8e4m3Mixin):
    # The group dimension for the weights it is automatically identified based on the layer type
    # If a new layer type is used, it can be manually specified
    pass

class MXFloat8Act(MXFloat8e4m3Act, Fp8e4m3Mixin):
    # In layerwise quantization, groupdim is automatically determined
    pass


class MXModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = qnn.QuantConv2d(32, 64, 3, weight_quant=MXFloat8Weight, input_quant=MXFloat8Act)

    def forward(self, x):
        return self.conv(x)

mx_model = MXModel()
x = torch.randn(1, 32, 8, 8)
mx_model.eval()
o = mx_model(x)

intermediate_input = mx_model.conv.input_quant(x)
assert isinstance(intermediate_input, GroupwiseFloatQuantTensor)
assert isinstance(mx_model.conv.quant_weight(), GroupwiseFloatQuantTensor)

/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/quant/solver/act.py:134: UserWarning: Group dim is being selected assuming batched input. Using unbatched input will fail and requires manually specification of group_dim
  warn(

If the input channel dimension is not divisible by group size, padding will be applied.

[3]:

class MXFloat8WeightNoPadding(MXFloat8e4m3Weight, Fp8e4m3Mixin):
    # The group dimension for the weights it is automatically identified based on the layer type
    # If a new layer type is used, it can be manually specified
    group_size = 8

class MXFloat8ActNoPadding(MXFloat8e4m3Act, Fp8e4m3Mixin):
    # In layerwise quantization, groupdim is automatically determined
    group_size = 8


class MXModelNoPadding(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = qnn.QuantConv2d(8, 64, 3, weight_quant=MXFloat8WeightNoPadding, input_quant=MXFloat8ActNoPadding)

    def forward(self, x):
        return self.conv(x)

class MXModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = qnn.QuantConv2d(8, 64, 3, weight_quant=MXFloat8Weight, input_quant=MXFloat8Act)

    def forward(self, x):
        return self.conv(x)

mx_model_no_padding = MXModelNoPadding()
mx_model = MXModel()
# Make sure that the modules are the same
mx_model_no_padding.load_state_dict(mx_model.state_dict())

x = torch.randn(1, 8, 8, 8)
mx_model.eval()
mx_model_no_padding.eval()
o_no_padding = mx_model_no_padding(x)
o = mx_model(x)

# The quant weight of the padded model is different from the non padding one
print(f"Non padding weights shape {mx_model_no_padding.conv.quant_weight().value_.shape}")
print(f"Padded weights shape {mx_model.conv.quant_weight().value_.shape}")

# However, results are still the same
assert torch.allclose(o, o_no_padding)

Non padding weights shape torch.Size([64, 1, 8, 3, 3])
Padded weights shape torch.Size([64, 1, 32, 3, 3])

[ ]:

from brevitas.quant.mx_quant_ocp import MXFloat8e4m3WeightMSE
from brevitas.quant_tensor import GroupwiseFloatQuantTensor


class MXFloat8Weight(MXFloat8e4m3WeightMSE, Fp8e4m3Mixin):
    # The group dimension for the weights it is automatically identified based on the layer type
    # If a new layer type is used, it can be manually specified
    pass

class MXFloat8Act(MXFloat8e4m3Act, Fp8e4m3Mixin):
    # In layerwise quantization, groupdim is automatically determined
    pass


class MXModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = qnn.QuantConv2d(32, 64, 3, weight_quant=MXFloat8Weight, input_quant=MXFloat8Act)

    def forward(self, x):
        return self.conv(x)

mx_model = MXModel()
x = torch.randn(1, 32, 8, 8)
mx_model.eval()
o = mx_model(x)

intermediate_input = mx_model.conv.input_quant(x)
assert isinstance(intermediate_input, GroupwiseFloatQuantTensor)

[ ]:

from brevitas.quant_tensor import GroupwiseIntQuantTensor
from brevitas.quant.mx_quant_ocp import MXInt8Weight
from brevitas.quant.mx_quant_ocp import MXInt8Act
import torch.nn as nn
import brevitas.nn as qnn
import torch

class MXInt8Weight(MXInt8Weight):
    # The group dimension for the weights it is automatically identified based on the layer type
    # If a new layer type is used, it can be manually specified
    pass

class MXInt8Act(MXInt8Act):
    # In layerwise quantization, groupdim is automatically determined
    pass

class MXModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = qnn.QuantConv2d(32, 64, 3, weight_quant=MXInt8Weight, input_quant=MXInt8Act)

    def forward(self, x):
        return self.conv(x)

mx_model = MXModel()
x = torch.randn(1, 32, 8, 8)
mx_model.eval()
o = mx_model(x)

intermediate_input = mx_model.conv.input_quant(x)
assert isinstance(intermediate_input, GroupwiseIntQuantTensor)
On this page
Show Source