added support fp8 types from OpenVINO

Author: alexsu52

nncf/openvino/graph/nncf_graph_builder.py

@@ -44,11 +44,13 @@ def convert_to_nncf_dtype(ov_type: ov.Type) -> Dtype:
         """
         type_name = ov_type.get_type_name()
         conversion_map = {
+            "nf4": "float",
+            "f8e4m3": "float",
+            "f8e5m2": "float",
             "f16": "float",
             "bf16": "float",
             "f32": "float",
             "f64": "float",
-            "nf4": "float",
             "i4": "int",
             "i8": "int",
             "i16": "int",

nncf/tensor/definitions.py

@@ -39,6 +39,9 @@ class TensorDataType(Enum):
     bfloat16 = auto()
     float32 = auto()
     float64 = auto()
+    f8e4m3 = auto()
+    f8e5m2 = auto()
+    nf4 = auto()
     int8 = auto()
     int32 = auto()
     int64 = auto()

nncf/tensor/functions/openvino_numeric.py

@@ -21,6 +21,9 @@
 from nncf.tensor.functions import numeric
 
 DTYPE_MAP: Dict[TensorDataType, ov.Type] = {
+    TensorDataType.nf4: ov.Type.nf4,
+    TensorDataType.f8e4m3: ov.Type.f8e4m3,
+    TensorDataType.f8e5m2: ov.Type.f8e5m2,
     TensorDataType.float16: ov.Type.f16,
     TensorDataType.bfloat16: ov.Type.bf16,
     TensorDataType.float32: ov.Type.f32,
@@ -48,12 +51,17 @@ def _(a: ov.Tensor) -> TensorBackend:
 
 @numeric.astype.register
 def _(a: ov.Tensor, dtype: TensorDataType) -> ov.Tensor:
-    if a.get_element_type() in [ov.Type.bf16, ov.Type.i4, ov.Type.u4] or dtype in [
+    ov_cast_types = [
         TensorDataType.bfloat16,
         TensorDataType.int4,
         TensorDataType.uint4,
-    ]:
-        # Cannot cast to/from bfloat16, uint4, int4 directly
+        TensorDataType.nf4,
+        TensorDataType.f8e4m3,
+        TensorDataType.f8e5m2,
+    ]
+    a_dtype = DTYPE_MAP_REV[a.get_element_type()]
+    if a_dtype in ov_cast_types or dtype in ov_cast_types:
+        # Cast using OpenVINO because the target or source dtype requires special handling
         return _astype_ov(a, dtype)
     return ov.Tensor(numeric.astype(a.data, dtype).data)
 
@@ -75,9 +83,16 @@ def _(a: ov.Tensor, shape: Union[int, Tuple[int, ...]]) -> ov.Tensor:
 
 @numeric.as_numpy_tensor.register
 def _(a: ov.Tensor) -> NDArray[Any]:
-    # Cannot convert bfloat16, uint4, int4 to numpy directly
+    # Cannot convert bfloat16, uint4, int4, nf4, f8e4m3, f8e5m2 to numpy directly
     a_dtype = DTYPE_MAP_REV[a.get_element_type()]
-    if a_dtype in [TensorDataType.bfloat16, TensorDataType.uint4, TensorDataType.int4]:
+    if a_dtype in [
+        TensorDataType.bfloat16,
+        TensorDataType.uint4,
+        TensorDataType.int4,
+        TensorDataType.nf4,
+        TensorDataType.f8e4m3,
+        TensorDataType.f8e5m2,
+    ]:
         dtype = TensorDataType.float32
         if a_dtype == TensorDataType.uint4:
             dtype = TensorDataType.uint8

tests/cross_fw/test_templates/template_test_nncf_tensor.py

@@ -1512,6 +1512,9 @@ def test_fn_zeros(self):
                 in [
                     TensorDataType.int4,
                     TensorDataType.uint4,
+                    TensorDataType.nf4,
+                    TensorDataType.f8e4m3,
+                    TensorDataType.f8e5m2,
                 ]
             ):
                 continue
@@ -1541,6 +1544,9 @@ def test_fn_eye(self, n, m, ref):
                 in [
                     TensorDataType.int4,
                     TensorDataType.uint4,
+                    TensorDataType.nf4,
+                    TensorDataType.f8e4m3,
+                    TensorDataType.f8e5m2,
                 ]
             ):
                 continue

tests/openvino/native/test_nncf_graph_builder.py

@@ -39,6 +39,8 @@ def test_compare_nncf_graph_synthetic_models(model_cls_to_test):
     "model,precision",
     [
         (FPModel(const_dtype=ov.Type.nf4), "nf4"),
+        (FPModel(const_dtype=ov.Type.f8e4m3), "f8e4m3"),
+        (FPModel(const_dtype=ov.Type.f8e5m2), "f8e5m2"),
     ],
 )
 def test_compare_nncf_graph_precision_synthetic_models(model, precision):
@@ -112,10 +114,12 @@ def _get_default_nncf_graph_edge(from_node, to_node, input_port_id, output_port_
 @pytest.mark.parametrize(
     "ov_type,expected_nncf_dtype",
     [
+        (ov.Type.nf4, Dtype.FLOAT),
+        (ov.Type.f8e4m3, Dtype.FLOAT),
+        (ov.Type.f8e5m2, Dtype.FLOAT),
         (ov.Type.f16, Dtype.FLOAT),
         (ov.Type.f32, Dtype.FLOAT),
         (ov.Type.f64, Dtype.FLOAT),
-        (ov.Type.nf4, Dtype.FLOAT),
         (ov.Type.i4, Dtype.INTEGER),
         (ov.Type.i8, Dtype.INTEGER),
         (ov.Type.i16, Dtype.INTEGER),
@@ -140,8 +144,6 @@ def test_convert_to_nncf_dtype_supported_types(ov_type: ov.Type, expected_nncf_d
     "ov_type",
     [
         ov.Type.undefined,
-        ov.Type.f8e4m3,
-        ov.Type.f8e5m2,
     ],
 )
 def test_convert_to_nncf_dtype_unsupported_types(ov_type: ov.Type):

tests/openvino/native/test_tensor.py

@@ -26,14 +26,23 @@
 class TestOVNNCFTensorOperators:
     @staticmethod
     def to_tensor(x, backend=TensorBackend.ov, dtype=TensorDataType.float32):
+        no_numpy_support_dtypes = [
+            TensorDataType.bfloat16,
+            TensorDataType.uint4,
+            TensorDataType.int4,
+            TensorDataType.nf4,
+            TensorDataType.f8e5m2,
+            TensorDataType.f8e4m3,
+        ]
+
         if backend == TensorBackend.ov:
-            if dtype in [TensorDataType.bfloat16, TensorDataType.uint4, TensorDataType.int4]:
+            if dtype in no_numpy_support_dtypes:
                 ov_const = opset.constant(x, dtype=DTYPE_MAP_OV[dtype])
                 return ov.Tensor(ov_const.data, ov_const.data.shape, DTYPE_MAP_OV[dtype])
             else:
                 return ov.Tensor(np.array(x, dtype=DTYPE_MAP_NP[dtype]))
         elif backend == TensorBackend.numpy:
-            if dtype in [TensorDataType.bfloat16, TensorDataType.uint4, TensorDataType.int4]:
+            if dtype in no_numpy_support_dtypes:
                 msg = f"Can't create NumPY tensor in dtype {dtype}"
                 raise ValueError(msg)
             return np.array(x, dtype=DTYPE_MAP_NP[dtype])