Adding FP8 weight export

src/brevitas/export/common/handler/base.py

@@ -4,6 +4,7 @@
 from abc import ABC
 from abc import abstractmethod
 import math
+from warnings import warn
 
 import torch
 from torch import Tensor
@@ -12,7 +13,8 @@
 from brevitas.function.ops import max_int
 from brevitas.function.ops import min_int
 
-__all__ = ['BaseHandler', 'BitWidthHandlerMixin', 'ZeroPointHandlerMixin']
+__all__ = [
+ 'BaseHandler', 'BitWidthHandlerMixin', 'ZeroPointHandlerMixin', 'FloatZeroPointHandlerMixin']
 
 
 class BaseHandler(Module, ABC):
@@ -38,6 +40,13 @@ def quant_axis(cls, scale):
  return None
 
 
+class FloatClipMixin(ABC):
+
+ @classmethod
+ def clip_symbolic_kwargs(cls, narrow, signed, exponent_bit_width, mantissa_bit_width):
+ return None
+
+
 class ClipMixin(ABC):
 
  @classmethod
@@ -112,6 +121,18 @@ def validate_neg_scalar_int_exponent(cls, scale: Tensor):
  return -cls.validate_scalar_int_exponent(scale)
 
 
+class FloatZeroPointHandlerMixin(ABC):
+
+ @classmethod
+ def zero_point_with_dtype(cls, signed, exponent_bit_width, mantissa_bit_width, zero_point):
+ if exponent_bit_width == 4 and mantissa_bit_width == 3:
+ return zero_point.type(torch.float8_e4m3fn)
+ elif exponent_bit_width == 5 and mantissa_bit_width == 2:
+ return zero_point.type(torch.float8_e5m2)
+ else:
+ return zero_point.type(torch.float32)
+
+
 class ZeroPointHandlerMixin(ABC):
 
  @classmethod

src/brevitas/export/onnx/manager.py

@@ -30,6 +30,43 @@
 from ..manager import ExportContext
 
 
+# workaround for fp8 not having many operators implemented
+class Fp8Workaround():
+
+ def __init__(self):
+ self.lib = None
+
+ def __enter__(self):
+ if torch_version >= version.parse('2.1.0'):
+ self.lib = torch.library.Library("aten", "IMPL")
+
+ def equal_cpu(self, other):
+ if (isinstance(self, Tensor) and
+ self.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)) or (
+ isinstance(other, Tensor) and
+ other.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)):
+ self = self.to(torch.float32)
+ other = other.to(torch.float32)
+ return torch.equal(self, other)
+ else:
+ res = True
+ if not isinstance(self, Tensor):
+ self = torch.tensor(self)
+ if not isinstance(other, Tensor):
+ other = torch.tensor(other)
+ if self.dim() > 0:
+ for x, y in zip(self.flatten(), other.flatten()):
+ res &= x == y
+ else:
+ res = self.item() == other.item()
+ return torch.tensor([res])
+
+ self.lib.impl("equal", equal_cpu, "CPU")
+
+ def __exit__(self, exc_type, exc_value, exc_traceback):
+ self.lib = None
+
+
 class ONNXBaseManager(BaseManager, ABC):
 
  model_transforms = []
@@ -127,7 +164,9 @@ def export_onnx(
  else:
  model_bytes = BytesIO()
  export_target = model_bytes
- torch.onnx.export(module, args, export_target, **onnx_export_kwargs)
+
+ with Fp8Workaround():
+ torch.onnx.export(module, args, export_target, **onnx_export_kwargs)
 
  # restore the model to previous properties
  module.apply(lambda m: _restore_act_caching_mode(m))

src/brevitas/export/onnx/standard/qcdq/handler.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from abc import ABC
+from warnings import warn
 
 import torch
 
@@ -10,6 +11,9 @@
 from brevitas.export.common.handler.qcdq import DQCastMixin
 from brevitas.export.common.handler.qcdq import DynamicQDQCastActQuantProxyHandlerMixin
 from brevitas.export.common.handler.qcdq import DynamicQMixin
+from brevitas.export.common.handler.qcdq import FloatQCDQCastActQuantProxyHandlerMixin
+from brevitas.export.common.handler.qcdq import FloatQCDQCastWeightQuantProxyHandlerMixin
+from brevitas.export.common.handler.qcdq import FloatQMixin
 from brevitas.export.common.handler.qcdq import QCDQCastActQuantProxyHandlerMixin
 from brevitas.export.common.handler.qcdq import QCDQCastDecoupledWeightQuantProxyHandlerMixin
 from brevitas.export.common.handler.qcdq import \
@@ -47,12 +51,33 @@ def validate(self, module):
  assert module.bit_width() > 1., 'Binary quant not supported'
 
 
+class StdFloatDQCastONNXMixin(StdDQCastONNXMixin, ABC):
+
+ def validate(self, module):
+ pass
+
+
+class StdFloatCDQCastONNXMixin(CDQCastMixin, StdFloatDQCastONNXMixin, ABC):
+
+ def clip_fn(self, x, min_val, max_val):
+ return IntClipFn.apply(x, min_val, max_val)
+
+
 class StdCDQCastONNXMixin(CDQCastMixin, StdDQCastONNXMixin, ABC):
 
  def clip_fn(self, x, min_val, max_val):
  return IntClipFn.apply(x, min_val, max_val)
 
 
+class StdFloatQCDQCastONNXMixin(FloatQMixin, StdFloatCDQCastONNXMixin, ABC):
+
+ def validate(self, module):
+ pass
+
+ def quantize_fn(self, x, scale, zero_point, dtype, axis):
+ return QuantizeLinearFn.apply(x, scale, zero_point, dtype, axis)
+
+
 class StdQCDQCastONNXMixin(QMixin, StdCDQCastONNXMixin, ABC):
 
  @classmethod
@@ -112,6 +137,12 @@ def quantize_fn(self, x, dtype):
  return DynamicQuantizeLinearFn.apply(x, dtype)
 
 
+class StdFloatQCDQCastONNXWeightQuantProxyHandler(StdFloatQCDQCastONNXMixin,
+ FloatQCDQCastWeightQuantProxyHandlerMixin,
+ ONNXBaseHandler):
+ _export_q_node = False
+
+
 class StdQCDQCastONNXWeightQuantProxyHandler(StdQCDQCastONNXMixin,
  QCDQCastWeightQuantProxyHandlerMixin,
  ONNXBaseHandler):
@@ -130,6 +161,12 @@ class StdQCDQCastONNXDecoupledWeightQuantWithInputProxyHandler(
  _export_q_node = False
 
 
+class StdFloatQCDQCastONNXActQuantProxyHandler(StdFloatQCDQCastONNXMixin,
+ FloatQCDQCastActQuantProxyHandlerMixin,
+ ONNXBaseHandler):
+ pass
+
+
 class StdQCDQCastONNXActQuantProxyHandler(StdQCDQCastONNXMixin,
  QCDQCastActQuantProxyHandlerMixin,
  ONNXBaseHandler):

src/brevitas/export/onnx/standard/qcdq/manager.py

@@ -17,6 +17,8 @@
 from ..manager import StdONNXBaseManager
 from .handler import StdCDQCastONNXBiasQuantProxyHandler
 from .handler import StdDynamicQDQCastONNXActQuantProxyHandler
+from .handler import StdFloatQCDQCastONNXActQuantProxyHandler
+from .handler import StdFloatQCDQCastONNXWeightQuantProxyHandler
 from .handler import StdQCDQCastONNXActQuantProxyHandler
 from .handler import StdQCDQCastONNXDecoupledWeightQuantProxyHandler
 from .handler import StdQCDQCastONNXDecoupledWeightQuantWithInputProxyHandler
@@ -36,8 +38,10 @@ class StdQCDQONNXManager(StdONNXBaseManager):
 
  handlers = [
  StdQCDQCastONNXWeightQuantProxyHandler,
+ StdFloatQCDQCastONNXWeightQuantProxyHandler,
  StdCDQCastONNXBiasQuantProxyHandler,
  StdQCDQCastONNXActQuantProxyHandler,
+ StdFloatQCDQCastONNXActQuantProxyHandler,
  StdQCDQCastONNXDecoupledWeightQuantProxyHandler,
  StdDynamicQDQCastONNXActQuantProxyHandler,
  StdQCDQCastONNXTruncQuantProxyHandler,

src/brevitas/proxy/__init__.py

@@ -1,6 +1,8 @@
 # Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
 
+from .float_parameter_quant import WeightFloatQuantProxyFromInjector
+from .float_runtime_quant import ActFloatQuantProxyFromInjector
 from .parameter_quant import BiasQuantProxyFromInjector
 from .parameter_quant import BiasQuantProxyFromInjectorBase
 from .parameter_quant import DecoupledWeightQuantProxyFromInjector

src/brevitas/proxy/float_runtime_quant.py

@@ -14,43 +14,28 @@
 class ActFloatQuantProxyFromInjector(ActQuantProxyFromInjectorBase):
 
  def scale(self, force_eval=True):
- if self.is_quant_enabled:
- current_status = self.training
- if force_eval:
- self.eval()
- out = self.__call__(self._zero_hw_sentinel())
- self.train(current_status)
- return out.scale
- elif self._cached_act is not None:
- return self._cached_act.scale
- elif self._cached_act is None:
- return None
+ return self.retrieve_attribute('scale', force_eval)
 
  def zero_point(self, force_eval=True):
- if self.is_quant_enabled:
- current_status = self.training
- if force_eval:
- self.eval()
- out = self.__call__(self._zero_hw_sentinel())
- self.train(current_status)
- return out.zero_point
- elif self._cached_act is not None:
- return self._cached_act.zero_point
- elif self._cached_act is None:
- return None
+ return self.retrieve_attribute('zero_point', force_eval)
 
- def bit_width(self, force_eval=True):
- if self.is_quant_enabled:
- current_status = self.training
- if force_eval:
- self.eval()
- out = self.__call__(self._zero_hw_sentinel())
- self.train(current_status)
- return out.bit_width
- elif self._cached_act is not None:
- return self._cached_act.bit_width
- elif self._cached_act is None:
- return None
+ def exponent_bit_width(self, force_eval=True):
+ return self.retrieve_attribute('exponent_bit_width', force_eval)
+
+ def mantissa_bit_width(self, force_eval=True):
+ return self.retrieve_attribute('mantissa_bit_width', force_eval)
+
+ def exponent_bias(self, force_eval=True):
+ return self.retrieve_attribute('exponent_bias', force_eval)
+
+ def saturating(self, force_eval=True):
+ return self.retrieve_attribute('saturating', force_eval)
+
+ def inf_values(self, force_eval=True):
+ return self.retrieve_attribute('inf_values', force_eval)
+
+ def nan_values(self, force_eval=True):
+ return self.retrieve_attribute('nan_values', force_eval)
 
  def forward(self, x: Union[Tensor, FloatQuantTensor]) -> Union[Tensor, FloatQuantTensor]:
  out = x
@@ -68,7 +53,8 @@ def forward(self, x: Union[Tensor, FloatQuantTensor]) -> Union[Tensor, FloatQuan
  y = self.fused_activation_quant_proxy(y)
  # If y is an empty FloatQuantTensor, we need to check if this is a passthrough proxy,
  # otherwise return a simple Tensor
- if isinstance(y, tuple) and not any(map(lambda f: f is None, y)):
+ # We exclude the last two values (inf_values and nan_values)
+ if isinstance(y, tuple) and not any(map(lambda f: f is None, y[:-2])):
  out = FloatQuantTensor(*y, signed=self.is_signed, training=self.training)
  elif self.is_passthrough_act: # preserve scale/zp/bit/sign even without output quant
  if isinstance(y, tuple):

src/brevitas/proxy/runtime_quant.py

@@ -22,6 +22,7 @@
  'ActQuantProxyProtocol',
  'AccQuantProxyProtocol',
  'ActQuantProxyFromInjector',
+ 'FloatActQuantProxyFromInjector',
  'TruncQuantProxyFromInjector',
  'ClampQuantProxyFromInjector']
 
@@ -95,6 +96,23 @@ def __init__(self, quant_layer, quant_injector):
  self.cache_inference_quant_act = False
  self.cache_quant_io_metadata_only = True
 
+ def internal_forward(self, force_eval):
+ current_status = self.training
+ if force_eval:
+ self.eval()
+ out = self.__call__(self._zero_hw_sentinel())
+ self.train(current_status)
+ return out
+
+ def retrieve_attribute(self, attribute, force_eval):
+ if self.is_quant_enabled:
+ out = self.internal_forward(force_eval)
+ return getattr(out, attribute)
+ elif self._cached_act is not None:
+ return getattr(self._cached_act, attribute)
+ elif self._cached_act is None:
+ return None
+
  @property
  def is_quant_enabled(self):
  return self._is_quant_enabled and not self.disable_quant
@@ -132,43 +150,13 @@ def init_tensor_quant(self):
 class ActQuantProxyFromInjector(ActQuantProxyFromInjectorBase):
 
  def scale(self, force_eval=True):
- if self.is_quant_enabled:
- current_status = self.training
- if force_eval:
- self.eval()
- out = self.__call__(self._zero_hw_sentinel())
- self.train(current_status)
- return out.scale
- elif self._cached_act is not None:
- return self._cached_act.scale
- elif self._cached_act is None:
- return None
+ return self.retrieve_attribute('scale', force_eval)
 
  def zero_point(self, force_eval=True):
- if self.is_quant_enabled:
- current_status = self.training
- if force_eval:
- self.eval()
- out = self.__call__(self._zero_hw_sentinel())
- self.train(current_status)
- return out.zero_point
- elif self._cached_act is not None:
- return self._cached_act.zero_point
- elif self._cached_act is None:
- return None
+ return self.retrieve_attribute('zero_point', force_eval)
 
  def bit_width(self, force_eval=True):
- if self.is_quant_enabled:
- current_status = self.training
- if force_eval:
- self.eval()
- out = self.__call__(self._zero_hw_sentinel())
- self.train(current_status)
- return out.bit_width
- elif self._cached_act is not None:
- return self._cached_act.bit_width
- elif self._cached_act is None:
- return None
+ return self.retrieve_attribute('bit_width', force_eval)
 
  def forward(self, x: Union[Tensor, IntQuantTensor]) -> Union[Tensor, IntQuantTensor]:
  out = x

src/brevitas/quant/experimental/float_base.py

@@ -7,10 +7,8 @@
 from brevitas.core.scaling.float_scaling import FloatScaling
 from brevitas.inject import ExtendedInjector
 from brevitas.inject import value
-from brevitas.proxy.float_parameter_quant import WeightFloatQuantProxyFromInjector
-from brevitas.proxy.float_runtime_quant import ActFloatQuantProxyFromInjector
-from brevitas.proxy.parameter_quant import WeightQuantProxyFromInjector
-from brevitas.proxy.runtime_quant import ActQuantProxyFromInjector
+from brevitas.proxy import ActFloatQuantProxyFromInjector
+from brevitas.proxy import WeightFloatQuantProxyFromInjector
 from brevitas.quant.solver import ActQuantSolver
 from brevitas.quant.solver import WeightQuantSolver
 from brevitas.quant.solver.common import SolveTensorQuantFloatToIntImplFromEnum