Better Bfloat16 support

src/brevitas/core/stats/stats_op.py

@@ -12,6 +12,8 @@
 from brevitas import config
 from brevitas.core.utils import StatelessBuffer
 from brevitas.function.ops import max_int
+# Use custom implementation of kthvalue as work around to (b)float16 kernel limitations
+from brevitas.utils.torch_utils import kthvalue
 
 from .stats_wrapper import SCALAR_SHAPE
 
@@ -64,15 +66,15 @@ def forward(self, x: Tensor):
         if self.stats_reduce_dim is None:
             # k is 1-indexed, so round away from zero
             k = int(math.floor(.01 * self.q * x.numel() + 0.5))
-            result = x.abs().view(-1).kthvalue(k).values
+            result = kthvalue(x.abs().view(-1), k)[0]
         else:
             # assuming x is two dimensional, get the other dimension
             assert len(x.size()) == 2, "Only 2-dim input is supported."
             other_dim = abs(self.stats_reduce_dim - 1)
             dim_slice = torch.narrow(x, dim=other_dim, start=0, length=1)
             # k is 1-indexed, so round away from zero
             k = int(math.floor(.01 * self.q * dim_slice.numel() + 0.5))
-            result = x.abs().kthvalue(k, dim=self.stats_reduce_dim, keepdim=self.keepdim).values
+            result = kthvalue(x.abs(), k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
         return result
 
 
@@ -97,15 +99,15 @@ def forward(self, x: Tensor) -> Tensor:
         if self.stats_reduce_dim is None:
             # k is 1-indexed, so round away from zero
             k = int(math.ceil(.01 * self.q * x.numel()))
-            result = x.view(-1).kthvalue(k).values
+            result = kthvalue(x.view(-1), k)[0]
         else:
             # assuming x is two dimensional, get the other dimension
             assert len(x.size()) == 2, "Only 2-dim input is supported."
             other_dim = abs(self.stats_reduce_dim - 1)
             dim_slice = torch.narrow(x, dim=other_dim, start=0, length=1)
             # k is 1-indexed, so round away from zero
             k = int(math.ceil(.01 * self.q * dim_slice.numel()))
-            result = x.kthvalue(k, dim=self.stats_reduce_dim, keepdim=self.keepdim).values
+            result = kthvalue(x, k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
         result = torch.clamp(result, max=self.zero())
         return result
 
@@ -134,8 +136,8 @@ def forward(self, x: Tensor) -> Tensor:
             low_k = int(math.ceil(.01 * self.low_q * x.numel()))
             # k is 1-indexed, so round away from zero
             high_k = int(math.floor(.01 * self.high_q * x.numel() + 0.5))
-            low_result = x.view(-1).kthvalue(low_k).values
-            high_result = x.view(-1).kthvalue(high_k).values
+            low_result = kthvalue(x.view(-1), low_k)[0]
+            high_result = kthvalue(x.view(-1), high_k)[0]
         else:
             # assuming x is two dimensional, get the other dimension
             assert len(x.size()) == 2, "Only 2-dim input is supported."
@@ -144,8 +146,8 @@ def forward(self, x: Tensor) -> Tensor:
             low_k = int(math.ceil(.01 * self.low_q * dim_slice.numel()))
             # k is 1-indexed, so round away from zero
             high_k = int(math.floor(.01 * self.high_q * dim_slice.numel() + 0.5))
-            low_result = x.kthvalue(low_k, dim=self.stats_reduce_dim, keepdim=self.keepdim).values
-            high_result = x.kthvalue(high_k, dim=self.stats_reduce_dim, keepdim=self.keepdim).values
+            low_result = kthvalue(x, low_k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
+            high_result = kthvalue(x, high_k, dim=self.stats_reduce_dim, keepdim=self.keepdim)[0]
         # We need to make sure the lower bound is not positive to align with zero-point statistics
         low_result = torch.clamp(low_result, max=self.zero())
         interval = high_result - low_result

src/brevitas/quant_tensor/__init__.py

@@ -15,6 +15,7 @@
 from .torch_handler import QUANT_TENSOR_FN_HANDLER
 
 IS_VALID_ATOL = 2e-1
+BFLOAT16_IS_VALID_ATOL = 0.5
 
 
 class QuantTensorBase(NamedTuple):
@@ -104,8 +105,15 @@ def is_not_none(self):
 
     @property
     def _pre_round_int_value(self):
-        int_value = self.value / self.scale
-        int_value = int_value + self.zero_point
+        value = self.value
+        scale = self.scale
+        zero_point = self.zero_point
+        if self.scale.dtype == torch.bfloat16:
+            value = self.value.type(torch.float32)
+            scale = self.scale.type(torch.float32)
+            zero_point = self.zero_point.type(torch.float32)
+        int_value = value / scale
+        int_value = int_value + zero_point
         return int_value
 
     @property
@@ -114,8 +122,9 @@ def is_valid(self):
             with torch.no_grad():
                 pre_round_int_value = self._pre_round_int_value
                 rounded_int_value = torch.round(pre_round_int_value)
-                is_int = torch.isclose(
-                    pre_round_int_value, rounded_int_value, atol=IS_VALID_ATOL).all()
+                max_abs_diff = torch.max(torch.abs(pre_round_int_value - rounded_int_value))
+                atol = BFLOAT16_IS_VALID_ATOL if self.value.dtype == torch.bfloat16 else IS_VALID_ATOL
+                is_int = max_abs_diff < atol
                 if self.bit_width >= 2:
                     if self.signed:
                         is_upper_b = (2.0 ** (self.bit_width - 1) - 1 >= rounded_int_value).all()
@@ -176,7 +185,12 @@ def int(self, float_datatype=False):
         if self.is_valid:
             int_value = round_ste(self._pre_round_int_value)
             if float_datatype:
-                return int_value
+                # Values at 8bit and lower can be represented exactly with float16 and bfloat16
+                # otherwise (e.g. Int16 bias), we upscale to float32
+                if self.bit_width <= 8.:
+                    return int_value.type(self.scale.dtype)
+                else:
+                    return int_value.type(torch.float32)
             else:
                 if self.bit_width <= 8. and self.signed_t.item():
                     return int_value.to(torch.int8)
@@ -301,6 +315,8 @@ def cat(tensors, dim, out=None):
 
     def __neg__(self):
         neg_value = (-self.int(float_datatype=True) - self.zero_point) * self.scale
+        # In case the dtype of self.int is different from the one of the scale
+        neg_value = neg_value.type(self.scale.dtype)
         if self.signed:
             return QuantTensor(
                 value=neg_value,
@@ -432,6 +448,8 @@ def __truediv__(self, other):
     def __abs__(self):
         if self.signed:
             abs_value = (torch.abs(self.int(float_datatype=True)) - self.zero_point) * self.scale
+            # In case the dtype of self.int is different from the one of the scale
+            abs_value = abs_value.type(self.scale.dtype)
             return QuantTensor(
                 value=abs_value,
                 scale=self.scale,

src/brevitas/utils/torch_utils.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import copy
+from typing import Optional, Tuple
 
 import torch
 from torch.nn import Sequential
@@ -46,3 +47,36 @@ def torch_partial_deepcopy(model):
         memo[id(p)] = copy.copy(p)  # Shallow copy of parameters
     model_copy = copy.deepcopy(model, memo)
     return model_copy
+
+
+def kthvalue(
+    x: torch.Tensor,
+    k: int,
+    dim: Optional[int] = None,
+    keepdim: bool = False,
+    out: Optional[Tuple[torch.Tensor, torch.LongTensor]] = None
+) -> Tuple[torch.Tensor, torch.LongTensor]:
+    # As of torch 2.1, there is no kthvalue implementation:
+    # - In CPU for float16
+    # - In GPU for bfloat16
+    # In these cases we cast to float32 and then go back to the original dtype
+    dtype = x.dtype
+    device = str(x.device)
+
+    # We do not support out as buffer for the output, since we cannot control its dtype
+    if out is not None:
+        raise RuntimeError("out argument for kthvalue not supported")
+
+    if (dtype == torch.float16 and 'cpu' in device) or \
+        (dtype == torch.bfloat16 and 'cuda' in device):
+        x = x.type(torch.float32)
+
+    # PyTorch specify None as default for `dim` but it breaks if we specifically pass None
+    if dim is not None:
+        x, indices = torch.kthvalue(x, k, dim=dim, keepdim=keepdim)
+    else:
+        x, indices = torch.kthvalue(x, k, keepdim=keepdim)
+
+    if x.dtype != dtype:
+        x = x.type(dtype)
+    return (x, indices)

src/brevitas_examples/imagenet_classification/ptq/ptq_evaluate.py

@@ -72,6 +72,8 @@
     metavar='ARCH',
     choices=model_names,
     help='model architecture: ' + ' | '.join(model_names) + ' (default: resnet18)')
+parser.add_argument(
+    '--dtype', default='float', choices=['float', 'bfloat16'], help='Data type to use')
 parser.add_argument(
     '--target-backend',
     default='fx',
@@ -215,6 +217,7 @@
     default=None,
     type=int,
     help='Accumulator Bit Width for GPFA2Q (default: None)')
+parser.add_argument('--onnx-opset-version', default=None, type=int, help='ONNX opset version')
 add_bool_arg(parser, 'gptq', default=False, help='GPTQ (default: disabled)')
 add_bool_arg(parser, 'gpfq', default=False, help='GPFQ (default: disabled)')
 add_bool_arg(parser, 'gpfa2q', default=False, help='GPFA2Q (default: disabled)')
@@ -226,11 +229,11 @@
 
 def main():
     args = parser.parse_args()
+    dtype = getattr(torch, args.dtype)
 
     random.seed(SEED)
     np.random.seed(SEED)
     torch.manual_seed(SEED)
-
     if args.act_quant_calibration_type == 'stats':
         act_quant_calib_config = str(args.act_quant_percentile) + 'stats'
     else:
@@ -312,14 +315,15 @@ def main():
 
     # Get the model from torchvision
     model = get_torchvision_model(args.model_name)
+    model = model.to(dtype)
 
     # Preprocess the model for quantization
     if args.target_backend == 'flexml':
         # flexml requires static shapes, pass a representative input in
         img_shape = model_config['center_crop_shape']
         model = preprocess_for_flexml_quantize(
             model,
-            torch.ones(1, 3, img_shape, img_shape),
+            torch.ones(1, 3, img_shape, img_shape, dtype=dtype),
             equalize_iters=args.graph_eq_iterations,
             equalize_merge_bias=args.graph_eq_merge_bias,
             merge_bn=not args.calibrate_bn)
@@ -339,6 +343,7 @@ def main():
     # Define the quantized model
     quant_model = quantize_model(
         model,
+        dtype=dtype,
         backend=args.target_backend,
         scale_factor_type=args.scale_factor_type,
         bias_bit_width=args.bias_bit_width,
@@ -405,7 +410,7 @@ def main():
 
     # Validate the quant_model on the validation dataloader
     print("Starting validation:")
-    validate(val_loader, quant_model)
+    validate(val_loader, quant_model, stable=dtype != torch.bfloat16)
 
     if args.export_onnx_qcdq or args.export_torch_qcdq:
         # Generate reference input tensor to drive the export process
@@ -418,7 +423,7 @@ def main():
         export_name = os.path.join(args.export_dir, config)
         if args.export_onnx_qcdq:
             export_name = export_name + '.onnx'
-            export_onnx_qcdq(model, ref_input, export_name)
+            export_onnx_qcdq(model, ref_input, export_name, opset_version=args.onnx_opset_version)
         if args.export_torch_qcdq:
             export_name = export_name + '.pt'
             export_torch_qcdq(model, ref_input, export_name)

src/brevitas_examples/imagenet_classification/utils.py

@@ -61,7 +61,7 @@ def accuracy(output, target, topk=(1,), stable=False):
         return res
 
 
-def validate(val_loader, model):
+def validate(val_loader, model, stable=True):
     """
     Run validation on the desired dataset
     """
@@ -82,7 +82,7 @@ def print_accuracy(top1, prefix=''):
 
             output = model(images)
             # measure accuracy
-            acc1, = accuracy(output, target, stable=True)
+            acc1, = accuracy(output, target, stable=stable)
             top1.update(acc1[0], images.size(0))
 
         print_accuracy(top1, 'Total:')

tests/brevitas/core/test_stats.py

@@ -8,6 +8,8 @@
 from brevitas.core.stats import AbsPercentile
 from brevitas.core.stats import NegativePercentileOrZero
 from brevitas.core.stats import PercentileInterval
+# Use custom implementation of kthvalue as work around to (b)float16 kernel limitations
+from brevitas.utils.torch_utils import kthvalue
 
 
 def test_abs_percentile_per_tensor():
@@ -35,10 +37,10 @@ def compute_percentile(self, x, low_q=None, high_q=None):
         low_p, high_p = None, None
         if low_q is not None:
             k = int(math.ceil(.01 * low_q * x.numel()))
-            low_p = x.view(-1).kthvalue(k).values
+            low_p = kthvalue(x.view(-1), k=k)[0]
         if high_q is not None:
             k = int(math.floor(.01 * high_q * x.numel() + 0.5))
-            high_p = x.view(-1).kthvalue(k).values
+            high_p = kthvalue(x.view(-1), k=k)[0]
         return low_p, high_p
 
     def test_negative_percentile(self):

tests/brevitas/graph/test_calibration.py

@@ -12,6 +12,8 @@
 from brevitas.graph.calibrate import calibration_mode
 import brevitas.nn as qnn
 from brevitas.quant import Int8ActPerTensorFixedPoint
+# Use custom implementation of kthvalue as work around to (b)float16 kernel limitations
+from brevitas.utils.torch_utils import kthvalue
 from tests.brevitas.hyp_helper import float_tensor_random_size_st
 
 IN_CH = 8
@@ -21,7 +23,7 @@
 
 def compute_quantile(x, q):
     k = int(math.floor(.01 * q * x.numel() + 0.5))
-    result = x.abs().view(-1).kthvalue(k).values
+    result = kthvalue(x.abs().view(-1), k=k)[0]
     return result