[FP8] support FP8 quantization

Author: Li Wei

mppq/data.py

@@ -34,7 +34,10 @@ class DataType(IntEnum):
     INT64 = TensorProto.INT64
     UINT64 = TensorProto.UINT64
 
+    FP8_E4M3FN = TensorProto.FLOAT8E4M3FN
+    FP8_E4M3FNUZ = TensorProto.FLOAT8E4M3FNUZ
     FP8_E5M2 = TensorProto.FLOAT8E5M2
+    FP8_E5M2FNUZ = TensorProto.FLOAT8E5M2FNUZ
     BF16 = TensorProto.BFLOAT16
     FP16 = TensorProto.FLOAT16
     FP32 = TensorProto.FLOAT
@@ -89,7 +92,10 @@ def from_torch(cls, dtype: torch_type):
             torch.float16: DataType.FP16,
             torch.float32: DataType.FP32,
             torch.float64: DataType.FP64,
+            torch.float8_e4m3fn: DataType.FP8_E4M3FN,
+            torch.float8_e4m3fnuz: DataType.FP8_E4M3FNUZ,
             torch.float8_e5m2: DataType.FP8_E5M2,
+            torch.float8_e5m2fnuz: DataType.FP8_E5M2FNUZ,
         }
         if dtype not in torch_converting_dict:
             raise TypeError(
@@ -129,7 +135,10 @@ def to_torch(cls, dtype) -> torch_type:
             DataType.FP16: torch.float16,
             DataType.FP32: torch.float32,
             DataType.FP64: torch.float64,
+            DataType.FP8_E4M3FN: torch.float8_e4m3fn,
+            DataType.FP8_E4M3FNUZ: torch.float8_e4m3fnuz,
             DataType.FP8_E5M2: torch.float8_e5m2,
+            DataType.FP8_E5M2FNUZ: torch.float8_e5m2fnuz,
         }
         assert isinstance(dtype, int)
         return torch_converting_dict[DataType(dtype)]
@@ -211,6 +220,10 @@ def convert_any_to_numpy(x: Any, accept_none: bool = True) -> None | np.ndarray:
             raise ValueError("Trying to convert an empty value.")
         return x
     elif isinstance(x, torch.Tensor):
+        if "float8" in str(x.dtype):
+            return convert_any_to_numpy(
+                x.cpu().to(torch.float32).numpy(), accept_none=accept_none
+            )
         return convert_any_to_numpy(x.cpu().numpy(), accept_none=accept_none)
     elif isinstance(x, Number):
         return np.array([x])

mppq/ffi.py

@@ -152,7 +152,7 @@ def dummy_locator():
 class ENABLE_CUDA_KERNEL:
     """Auto config compiler path before entering compiling CUDA context"""
 
-    USING_CUDA_KERNEL = False
+    USING_CUDA_KERNEL = True
 
     def __init__(self) -> None:
         self._state = True

mppq/frontend/onnx/onnxruntime_exporter.py

@@ -15,7 +15,7 @@
     QuantVisibility,
     TensorQuantizationConfig,
 )
-from mppq.utils.qfunction import ppq_quant_toint
+from mppq.utils.qfunction import ppq_quant_toint, ppq_quant_tofloat
 from mppq.utils.round import ppq_tensor_round
 
 
@@ -91,6 +91,12 @@ def infer_qtype(self, config: TensorQuantizationConfig):
         if config.num_of_bits > 8:
             offset_dtype = torch.int32
             value_dtype = torch.int32
+        if config.exponent_bits == 4:
+            offset_dtype = torch.float8_e4m3fn
+            value_dtype = torch.float8_e4m3fn
+        if config.exponent_bits == 5:
+            offset_dtype = torch.float8_e5m2
+            value_dtype = torch.float8_e5m2
         return offset_dtype, value_dtype
 
     def insert_quantize_node(
@@ -141,17 +147,18 @@ def insert_quantize_node(
         elif config.policy.has_property(QuantizationProperty.FLOATING):
             # Following code will export Linear Quantization Config
             # That is for FP32 -> FP8
+            offset_dtype, value_type = self.infer_qtype(config)
             scale = convert_any_to_tensor(config.scale.clone(), dtype=torch.float32)
-            offset = convert_any_to_tensor(config.offset.clone(), dtype=torch.float32)
+            offset = convert_any_to_tensor(config.offset.clone(), dtype=offset_dtype)
 
             created = graph.create_operation(
-                op_type="QuantizeFloating",
-                attributes={
-                    "min": config.quant_min,
-                    "max": config.quant_max,
-                    "exponent": config.exponent_bits,
-                    "mantissa": config.mantissa_bits,
-                },
+                op_type="QuantizeLinear",
+                # attributes={
+                #     "min": config.quant_min,
+                #     "max": config.quant_max,
+                #     "exponent": config.exponent_bits,
+                #     "mantissa": config.mantissa_bits,
+                # },
             )
 
             if config.policy.has_property(QuantizationProperty.PER_CHANNEL):
@@ -171,10 +178,11 @@ def insert_quantize_node(
             graph.create_variable(
                 name=None, value=scale, is_parameter=True, dest_ops=[created]
             )
-            graph.create_variable(
-                name=None, value=offset, is_parameter=True, dest_ops=[created]
-            )
+            # graph.create_variable(
+            #     name=None, value=offset, is_parameter=True, dest_ops=[created]
+            # ) # zero_point is not used for FP8
 
+            created.outputs[0].dtype = value_type
             created.outputs[0].shape = var.shape
             created.inputs[0].shape = var.shape
             return created
@@ -231,17 +239,18 @@ def insert_dequantize_node(
             return created
 
         elif config.policy.has_property(QuantizationProperty.FLOATING):
+            offset_dtype, value_type = self.infer_qtype(config)
             scale = convert_any_to_tensor(config.scale.clone(), dtype=torch.float32)
-            offset = convert_any_to_tensor(config.offset.clone(), dtype=torch.float32)
+            offset = convert_any_to_tensor(config.offset.clone(), dtype=offset_dtype)
 
             created = graph.create_operation(
-                op_type="DequantizeFloating",
-                attributes={
-                    "min": config.quant_min,
-                    "max": config.quant_max,
-                    "exponent": config.exponent_bits,
-                    "mantissa": config.mantissa_bits,
-                },
+                op_type="DequantizeLinear",
+                # attributes={
+                #     "min": config.quant_min,
+                #     "max": config.quant_max,
+                #     "exponent": config.exponent_bits,
+                #     "mantissa": config.mantissa_bits,
+                # },
             )
 
             if config.policy.has_property(QuantizationProperty.PER_CHANNEL):
@@ -261,12 +270,14 @@ def insert_dequantize_node(
             graph.create_variable(
                 name=None, value=scale, is_parameter=True, dest_ops=[created]
             )
-            graph.create_variable(
-                name=None, value=offset, is_parameter=True, dest_ops=[created]
-            )
+            # graph.create_variable(
+            #     name=None, value=offset, is_parameter=True, dest_ops=[created]
+            # )
 
-            created.outputs[0].shape = var.shape
             created.inputs[0].shape = var.shape
+            created.inputs[0].dtype = value_type
+            created.outputs[0].shape = var.shape
+            created.outputs[0].dtype = torch.float32
 
             return created
         else:
@@ -468,6 +479,11 @@ def convert_operation(
                 ):
                     var.value = ppq_quant_toint(tensor=var.value, config=config)
 
+                if quantized_param and config.policy.has_property(
+                    QuantizationProperty.FLOATING
+                ):
+                    var.value = ppq_quant_tofloat(tensor=var.value, config=config)
+
             elif not var.is_parameter:
 
                 # Patch 20230103:

mppq/utils/qfunction/__init__.py

@@ -6,7 +6,7 @@
     QuantizationStates,
     TensorQuantizationConfig,
 )
-from mppq.utils.qfunction.floating import floating_quant
+from mppq.utils.qfunction.floating import floating_quant, floating_quant_tofloat
 from mppq.utils.qfunction.linear import (
     dynamic_linear_quant,
     linear_fake_quant,
@@ -80,8 +80,29 @@ def ppq_quant_toint(
     )
 
 
+def ppq_quant_tofloat(
+    tensor: torch.Tensor, config: TensorQuantizationConfig
+) -> torch.Tensor:
+    """
+    ## PPQ 核心量化函数
+
+    根据 config 中描述的策略，这个函数将会执行线性量化，动态量化
+
+    但是结果直接出来是float8类型
+    """
+    if config.policy.has_property(QuantizationProperty.FLOATING):
+        if not config.policy.has_property(QuantizationProperty.DYNAMIC):
+            return floating_quant_tofloat(tensor, config)
+
+    raise ValueError(
+        "Unexpected Quantization Property Found in ppq_quant_tofp8. "
+        "Do not know how to quantize your config yet."
+    )
+
+
 __all__ = [
     "ppq_fake_quant",
     "ppq_quant_toint",
+    "ppq_quant_tofloat",
     "BaseQuantFunction",
 ]

mppq/utils/qfunction/floating.py

@@ -137,3 +137,25 @@ def floating_quant(
         )
     assert isinstance(qtensor, torch.Tensor)
     return qtensor
+
+
+def floating_quant_tofloat(
+    tensor: torch.Tensor, config: TensorQuantizationConfig
+) -> torch.Tensor:
+    """PPQ 核心量化函数，没啥好说的了吧，这个玩意只做 quant 不做 dequant"""
+    if not config.policy.has_property(QuantizationProperty.FLOATING):
+        raise ValueError("Critical Quantization Error! Non-floating config detected.")
+
+    if config.policy.has_property(QuantizationProperty.PER_CHANNEL):
+        shape = [
+            1 if axis != config.channel_axis else -1 for axis in range(tensor.ndim)
+        ]
+        scale = config.scale.view(shape)
+        offset = config.offset.view(shape).to(tensor.device)
+        tensor = (tensor / scale) + offset
+    else:  # QuantizationProperty.PER_TENSOR
+        tensor = (tensor / config.scale.to(tensor.device)) + config.offset.to(
+            tensor.device
+        )
+    tensor = torch.clamp(tensor, config.quant_min, config.quant_max)
+    return tensor.to(torch.float8_e4m3fn)