[PT FE] Support F8 constants (#29313)

### Details:
 - *Support F8 constants*
 - *Allow to set patch condition externally*

### Tickets:
 - *ticket-id*

Signed-off-by: Maxim Vafin <[email protected]>

Author: mvafin

src/bindings/python/src/openvino/frontend/pytorch/patch_model.py

@@ -76,25 +76,30 @@ def unpatch_model(model, orig_forward_name):
                             "Original exception details:\n%s", error)
 
 
-def __make_16bit_traceable(model: torch.nn.Module):
+def __make_16bit_traceable(model: torch.nn.Module,
+                           orig_forward_name: str = "_openvino_module_extension_patch_orig_forward",
+                           patch_condition=None):
     """
     Prepare a 16-bit PyTorch model for tracing with OpenVINO.
      - Replace known list of modules with ModuleExtension.
      - Convert other modules with weights to FP32.
     """
-    def patch_condition(module):
-        supported = [torch.float32, torch.float16, torch.bfloat16]
-        return (hasattr(module, "weight")
-                and getattr(module.weight, "dtype", None) in supported)
-
+    if patch_condition is None:
+        def patch_condition(module):
+            supported = {torch.float32, torch.float16, torch.bfloat16}
+            weight = getattr(module, "weight", None)
+            return weight is not None and weight.dtype in supported
+
+    def fp32_tensor(*shape):
+        return torch.full(shape, 0.5, dtype=torch.float32)
+    
     extensions = {
         torch.nn.Linear: ModuleExtension(
             torch.nn.Linear, "ov_ext::linear",
             convert=lambda module, target_op, *args, **kwargs: target_op(args[0],
                                                                          module.weight,
                                                                          module.bias),
-            evaluate=lambda module, *args, **kwargs: torch.full(
-                list(args[0].shape[:-1]) + [module.out_features], 0.5, dtype=torch.float32),
+            evaluate=lambda module, *args, **kwargs: fp32_tensor(*args[0].shape[:-1], module.out_features),
             condition=patch_condition),
         torch.nn.Embedding: ModuleExtension(
             torch.nn.Embedding, "ov_ext::embedding",
@@ -103,8 +108,7 @@ def patch_condition(module):
                                                                          module.padding_idx,
                                                                          module.scale_grad_by_freq,
                                                                          module.sparse),
-            evaluate=lambda module, *args, **kwargs: torch.full(
-                list(args[1].shape) + [module.embedding_dim], 0.5, dtype=torch.float32),
+            evaluate=lambda module, *args, **kwargs: fp32_tensor(*args[1].shape, module.embedding_dim),
             condition=patch_condition),
     }
     try:
@@ -114,14 +118,12 @@ def patch_condition(module):
             convert=lambda module, target_op, *args, **kwargs: target_op(args[0],
                                                                          module.weight,
                                                                          module.bias),
-            evaluate=lambda module, *args, **kwargs: torch.full(
-                list(args[0].shape[:-1]) + [module.nf], 0.5, dtype=torch.float32),
+            evaluate=lambda module, *args, **kwargs: fp32_tensor(*args[0].shape[:-1], module.nf),
             condition=patch_condition)
     except ImportError:
         pass
-    patch_model(model, extensions,
-                "_openvino_module_extension_patch_orig_forward")
-    dtype_to_patch = [torch.float16, torch.bfloat16]
+    patch_model(model, extensions, orig_forward_name)
+    dtype_to_patch = {torch.float16, torch.bfloat16}
     for _, module in model.named_modules():
         if (module.__class__ not in extensions and
             (any(p.dtype in dtype_to_patch for p in module.parameters(False))

src/bindings/python/src/openvino/frontend/pytorch/utils.py

@@ -49,23 +49,35 @@ def get_type_from_py_type(value):
     return OVType.dynamic
 
 
+F8_DTYPE_MAP = {
+    torch.float8_e4m3fn: OVType.f8e4m3,
+    torch.float8_e5m2: OVType.f8e5m2,
+}
+
+
 def torch_tensor_to_ov_const(torch_t: torch.Tensor, shared_memory=True):
-    is_fake_tensor = False
     try:
         from torch._prims import FakeTensor
-        is_fake_tensor = isinstance(torch_t, FakeTensor)
-    except:
+        if isinstance(torch_t, FakeTensor):
+            raise AssertionError("`FakeTensor` detected. Infer the "
+                                 "model before exporting to avoid this.")
+    except ImportError:
         pass
-    assert not is_fake_tensor, '`FakeTensor` is found in the graph during conversion. ' \
-                               'In order to avoid `FakeTensor` in the traced model, ' \
-                               'try to infer the model before exporting.'
+
+    dtype = torch_t.dtype
     torch_t = torch_t.contiguous()
-    if torch_t.dtype == torch.bfloat16:
+    if dtype == torch.bfloat16:
         # reinterpret bfloat16 data as float16 to allow conversion to numpy
         torch_t = torch_t.view(torch.float16)
         narr = torch_t.numpy(force=True)
         tensor = Tensor(narr, torch_t.shape, OVType.bf16)
         ov_const = op.Constant(tensor, shared_memory=shared_memory)
+    elif dtype in F8_DTYPE_MAP:
+        # reinterpret f8 data as u8 to allow conversion to numpy
+        torch_t = torch_t.view(torch.uint8)
+        narr = torch_t.numpy(force=True)
+        tensor = Tensor(narr, torch_t.shape, F8_DTYPE_MAP[dtype])
+        ov_const = op.Constant(tensor, shared_memory=shared_memory)
     else:
         narr = torch_t.numpy(force=True)
         ov_const = op.Constant(narr, shared_memory=shared_memory)
@@ -126,6 +138,8 @@ def graph_has_ops(graph, op_types: list) -> bool:
     "float": OVType.f32,
     "int": OVType.i64,
     "bool": OVType.boolean,
+    "torch.float8_e4m3fn": OVType.f8e4m3,
+    "torch.float8_e5m2": OVType.f8e5m2,
     "torch.bfloat16": OVType.bf16,
     "torch.float16": OVType.f16,
     "torch.float32": OVType.f32,

tests/layer_tests/py_frontend_tests/test_torch_decoder.py

@@ -96,10 +96,54 @@ def test_pytorch_decoder_get_input_type_none():
     assert isinstance(nc_decoder.get_input_type(2).value, DecoderType.PyNone)
 
 
+@pytest.mark.precommit
+def test_pytorch_decoder_can_convert_f8_e4m3_tensor():
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+    from openvino import PartialShape, Type
+
+    class SomeTensor(torch.nn.Module):
+        def forward(self):
+            return torch.tensor([1, 2], dtype=torch.float8_e4m3fn)
+
+    model = get_scripted_model(SomeTensor())
+    consts = [n for n in model.inlined_graph.nodes() if n.kind() ==
+              "prim::Constant"]
+    assert len(consts) > 0
+    some_const = consts[0]
+    nc_decoder = TorchScriptPythonDecoder(model, some_const)
+    ov_const = nc_decoder.as_constant()
+    assert ov_const is not None
+    assert len(ov_const) == 1
+    assert ov_const[0].get_element_type() == Type.f8e4m3
+    assert ov_const[0].get_partial_shape() == PartialShape([2])
+
+
+@pytest.mark.precommit
+def test_pytorch_decoder_can_convert_f8_e5m2_tensor():
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+    from openvino import PartialShape, Type
+
+    class SomeTensor(torch.nn.Module):
+        def forward(self):
+            return torch.tensor([1, 2], dtype=torch.float8_e5m2)
+
+    model = get_scripted_model(SomeTensor())
+    consts = [n for n in model.inlined_graph.nodes() if n.kind() ==
+              "prim::Constant"]
+    assert len(consts) > 0
+    some_const = consts[0]
+    nc_decoder = TorchScriptPythonDecoder(model, some_const)
+    ov_const = nc_decoder.as_constant()
+    assert ov_const is not None
+    assert len(ov_const) == 1
+    assert ov_const[0].get_element_type() == Type.f8e5m2
+    assert ov_const[0].get_partial_shape() == PartialShape([2])
+
+
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_fp16_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -121,7 +165,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bf16_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -143,7 +187,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_fp32_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -165,7 +209,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_fp64_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -187,7 +231,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bool_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -209,7 +253,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_u8_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -231,7 +275,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i8_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -253,7 +297,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i16_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -275,7 +319,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i32_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -297,7 +341,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i64_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def forward(self):
@@ -337,7 +381,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_int_list():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class ListConst(torch.nn.Module):
         def forward(self):
@@ -360,7 +404,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_float_list():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class ListConst(torch.nn.Module):
         def forward(self):
@@ -383,7 +427,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bool_list():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class ListConst(torch.nn.Module):
         def forward(self):
@@ -406,7 +450,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_int_tuple():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class ListConst(torch.nn.Module):
         def forward(self):
@@ -429,7 +473,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_float_tuple():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class ListConst(torch.nn.Module):
         def forward(self):
@@ -452,7 +496,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bool_tuple():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class ListConst(torch.nn.Module):
         def forward(self):
@@ -475,7 +519,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_empty_list():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class aten_roll(torch.nn.Module):
         def __init__(self, shifts):
@@ -503,7 +547,7 @@ def forward(self, x):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_int_scalar_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def __init__(self) -> None:
@@ -534,7 +578,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_float_scalar_tensor():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
         def __init__(self) -> None:
@@ -565,7 +609,7 @@ def forward(self):
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_tensor_list():
     from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
-    from openvino.runtime import PartialShape, Type
+    from openvino import PartialShape, Type
     from typing import List, Optional
 
     class SomeTensor(torch.nn.Module):