[CPU] FullyConnected weights compression: mxfp4 (wei=f4e2m1, scales=f…

…8e8m0) support
dmitry-gorokhov · Jul 31, 2024 · 6d31b31 · 6d31b31
1 parent 36eebc2
commit 6d31b31
Show file tree

Hide file tree

Showing 21 changed files with 259 additions and 104 deletions.
diff --git a/...on/transformations/include/transformations/low_precision/mark_dequantization_subgraph.hpp b/...on/transformations/include/transformations/low_precision/mark_dequantization_subgraph.hpp
@@ -22,7 +22,9 @@ namespace pass {
 class TRANSFORMATIONS_API MarkDequantizationSubgraph : public MatcherPass {
 public:
     OPENVINO_RTTI("MarkDequantizationSubgraph", "0");
-    MarkDequantizationSubgraph(const element::TypeVector& precisions, const bool fold_subtract_const = false);
+    MarkDequantizationSubgraph(const element::TypeVector& precisions,
+                               const bool fold_subtract_const = false,
+                               const bool fold_multiply_const = true);
 };
 }  // namespace pass
 }  // namespace ov
diff --git a/...common/transformations/src/transformations/low_precision/mark_dequantization_subgraph.cpp b/...common/transformations/src/transformations/low_precision/mark_dequantization_subgraph.cpp
@@ -14,7 +14,8 @@
 #include "transformations/utils/utils.hpp"
 
 ov::pass::MarkDequantizationSubgraph::MarkDequantizationSubgraph(const element::TypeVector& precisions,
-                                                                 const bool fold_subtract_const) {
+                                                                 const bool fold_subtract_const,
+                                                                 const bool fold_multiply_const) {
     // Dequantization subgraph may have two forms: with and without Subtract
     //
     //    Input                                 Input
@@ -99,6 +100,18 @@ ov::pass::MarkDequantizationSubgraph::MarkDequantizationSubgraph(const element::
 
         // mark Multiply as dequantization node
         ov::mark_as_dequantization_node(multiply);
+        auto scale = multiply->get_input_node_shared_ptr(1);
+        if (ov::is_type<ov::op::v0::Convert>(scale) &&
+            ov::is_type<ov::op::v0::Constant>(scale->get_input_node_ptr(0))) {
+            if (!fold_multiply_const) {
+                ov::disable_constant_folding(scale);
+                ov::unmark_as_decompression(scale);
+                ov::enable_keep_const_precision(scale->get_input_node_shared_ptr(0));
+            } else {
+                ov::enable_constant_folding(scale);
+                ov::disable_keep_const_precision(scale->get_input_node_shared_ptr(0));
+            }
+        }
 
         return false;
     };

diff --git a/src/frontends/onnx/tests/tests_python/test_backend.py b/src/frontends/onnx/tests/tests_python/test_backend.py
@@ -510,13 +510,9 @@ def expect_fail(test_case_path, xfail):  # type: (str) -> None
         "OnnxBackendNodeModelTest.test_cast_FLOAT_to_FLOAT8E5M2FNUZ_cpu",
         "OnnxBackendNodeModelTest.test_cast_FLOAT_to_FLOAT8E5M2_cpu",
         "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT16_to_FLOAT8E4M3FNUZ_cpu",
-        "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT16_to_FLOAT8E4M3FN_cpu",
         "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT16_to_FLOAT8E5M2FNUZ_cpu",
-        "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT16_to_FLOAT8E5M2_cpu",
         "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT_to_FLOAT8E4M3FNUZ_cpu",
-        "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT_to_FLOAT8E4M3FN_cpu",
         "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT_to_FLOAT8E5M2FNUZ_cpu",
-        "OnnxBackendNodeModelTest.test_cast_no_saturate_FLOAT_to_FLOAT8E5M2_cpu",
         "OnnxBackendNodeModelTest.test_castlike_FLOAT8E4M3FNUZ_to_FLOAT_cpu",
         "OnnxBackendNodeModelTest.test_castlike_FLOAT8E4M3FNUZ_to_FLOAT_expanded_cpu",
         "OnnxBackendNodeModelTest.test_castlike_FLOAT8E4M3FN_to_FLOAT_cpu",

diff --git a/src/plugins/intel_cpu/src/dnnl_extension_utils.cpp b/src/plugins/intel_cpu/src/dnnl_extension_utils.cpp
@@ -35,6 +35,8 @@ uint8_t DnnlExtensionUtils::sizeOfDataType(dnnl::memory::data_type dataType) {
     case dnnl::memory::data_type::nf4:
     case dnnl::memory::data_type::s4:
     case dnnl::memory::data_type::u4:
+    case dnnl::memory::data_type::f8_e8m0:
+    case dnnl::memory::data_type::f4_e2m1:
         return 1;
     case dnnl::memory::data_type::undef:
         return 0;
@@ -66,6 +68,10 @@ dnnl::memory::data_type DnnlExtensionUtils::ElementTypeToDataType(const ov::elem
             return memory::data_type::s4;
         case ov::element::u4:
             return memory::data_type::u4;
+        case ov::element::f8e8m0:
+            return memory::data_type::f8_e8m0;
+        case ov::element::f4e2m1:
+            return memory::data_type::f4_e2m1;
         case ov::element::undefined:
             return memory::data_type::undef;
         default: {
@@ -98,6 +104,10 @@ ov::element::Type DnnlExtensionUtils::DataTypeToElementType(const dnnl::memory::
             return ov::element::i4;
         case memory::data_type::u4:
             return ov::element::u4;
+        case memory::data_type::f8_e8m0:
+            return ov::element::f8e8m0;
+        case memory::data_type::f4_e2m1:
+            return ov::element::f4e2m1;
         case memory::data_type::undef:
             return ov::element::undefined;
         default: {

diff --git a/src/plugins/intel_cpu/src/dnnl_postops_composer.cpp b/src/plugins/intel_cpu/src/dnnl_postops_composer.cpp
@@ -615,12 +615,13 @@ static MemoryPtr prepackDecompressionParams(const MemoryCPtr& paramsPtr,
     return dstMem;
 }
 
-void DnnlPostOpsComposer::appendDecompressionScales(const MemoryCPtr& scales_ptr, bool needTranspose) {
+void DnnlPostOpsComposer::appendDecompressionScales(const MemoryCPtr& scales_ptr, bool needTranspose, ov::element::Type dstPrecision) {
     if (scales_ptr == nullptr)
         return;
 
-    auto scalesMem = prepackDecompressionParams(scales_ptr, needTranspose, ov::element::f32, engine);
-    attr.set_scales_dims(DNNL_ARG_WEIGHTS, DnnlExtensionUtils::convertToDnnlDims(scalesMem->getStaticDims()));
+    auto scalesMem = prepackDecompressionParams(scales_ptr, needTranspose, dstPrecision, engine);
+    attr.set_scales_dims(DNNL_ARG_WEIGHTS,
+        DnnlExtensionUtils::convertToDnnlDims(scalesMem->getStaticDims()), DnnlExtensionUtils::ElementTypeToDataType(dstPrecision));
     cpuArgs[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] = std::move(scalesMem);
     dnnlArgs[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] =
         cpuArgs[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS]->getPrimitive();

diff --git a/src/plugins/intel_cpu/src/dnnl_postops_composer.h b/src/plugins/intel_cpu/src/dnnl_postops_composer.h
@@ -31,7 +31,7 @@ class DnnlPostOpsComposer {
                         const bool hasBias,
                         const dnnl::memory::data_type outDataType);
     DnnlPrimitiveAttrs compose();
-    void appendDecompressionScales(const MemoryCPtr& scales_ptr, bool needTranspose);
+    void appendDecompressionScales(const MemoryCPtr& scales_ptr, bool needTranspose, ov::element::Type dstPrecision);
     void appendDecompressionZeroPoints(const MemoryCPtr& zero_points_ptr, bool needTranspose, ov::element::Type dstPrecision);
     void setDynamicQuantizationParams(uint64_t groupSize);
 

diff --git a/src/plugins/intel_cpu/src/graph_optimizer.cpp b/src/plugins/intel_cpu/src/graph_optimizer.cpp
@@ -289,7 +289,8 @@ void GraphOptimizer::FuseConvMatmulFCDeconvAndDQScales(Graph &graph) {
 }
 
 void GraphOptimizer::FuseFCAndWeightsDecompression(Graph &graph) {
-    std::set<ov::element::Type> supportedWeightsPrecisions{ov::element::u8, ov::element::i8, ov::element::nf4, ov::element::u4, ov::element::i4};
+    std::set<ov::element::Type> supportedWeightsPrecisions{
+        ov::element::u8, ov::element::i8, ov::element::nf4, ov::element::u4, ov::element::i4, ov::element::f4e2m1};
     const std::set<ov::element::Type> supportedDataPrecisions{ov::element::f32, ov::element::bf16};
     auto expectedNode = [](NodePtr node, Type expectedType) {
         return node->getType() == expectedType && node->getChildEdges().size() == 1;
@@ -329,16 +330,24 @@ void GraphOptimizer::FuseFCAndWeightsDecompression(Graph &graph) {
         }
 
         CPU_GRAPH_OPTIMIZER_SCOPE(FuseFCAndWeightsDecompression);
-        const auto multiplyConstNode = multiplyNode->getParentEdgeAt(1)->getParent();
+        const auto mulParent1 = multiplyNode->getParentEdgeAt(1)->getParent();
+        NodePtr multiplyParent, multiplyConvertNode, multiplyConstNode;
+        multiplyParent = mulParent1;
+        if (multiplyParent->getType() == Type::Convert) {
+            multiplyConvertNode = multiplyParent;
+            multiplyParent = multiplyConvertNode->getParentEdgeAt(0)->getParent();
+        }
+        multiplyConstNode = multiplyParent;
         if (multiplyConstNode->getType() != Type::Input) {
             SKIP_FUSION_FOR_NODE(fcNode);
         }
+        const bool withMultiplyConvert = multiplyConvertNode != nullptr;
 
-        const auto mulParent = multiplyNode->getParentEdgeAt(0)->getParent();
-        const bool withSubtract = mulParent->getAlgorithm() == Algorithm::EltwiseSubtract;
+        const auto mulParent0 = multiplyNode->getParentEdgeAt(0)->getParent();
+        const bool withSubtract = mulParent0->getAlgorithm() == Algorithm::EltwiseSubtract;
         NodePtr subtractNode, subtractConvertNode, subtractConstNode;
         if (withSubtract) {
-            subtractNode = mulParent;
+            subtractNode = mulParent0;
             if (!expectedNode(subtractNode, Type::Eltwise)) {
                 SKIP_FUSION_FOR_NODE(fcNode);
             }
@@ -354,7 +363,7 @@ void GraphOptimizer::FuseFCAndWeightsDecompression(Graph &graph) {
         }
 
         const bool withSubtractConvert = subtractConvertNode != nullptr;
-        const auto convertNode = withSubtract ? subtractNode->getParentEdgeAt(0)->getParent() : mulParent;
+        const auto convertNode = withSubtract ? subtractNode->getParentEdgeAt(0)->getParent() : mulParent0;
         if (!expectedNode(convertNode, Type::Convert)) {
             SKIP_FUSION_FOR_NODE(fcNode);
         }
@@ -461,6 +470,8 @@ void GraphOptimizer::FuseFCAndWeightsDecompression(Graph &graph) {
             fcNode->addOriginalLayer(subtractNode->getOriginalLayers());
         if (withSubtractConvert)
             fcNode->addOriginalLayer(subtractConvertNode->getOriginalLayers());
+        if (withMultiplyConvert)
+            fcNode->addOriginalLayer(multiplyConvertNode->getOriginalLayers());
 
         const auto& weightsPrecision = weightsNode->getOriginalOutputPrecisionAtPort(0);
         if (withTranspose) {
@@ -511,6 +522,12 @@ void GraphOptimizer::FuseFCAndWeightsDecompression(Graph &graph) {
                     graph.RemoveEdge(subtractConvertNode->getParentEdgeAt(0));
             }
             graph.RemoveEdge(multiplyNode->getParentEdgeAt(1));
+            if (withMultiplyConvert) {
+                // MultiplyConvert is removed only if there are no other consumers (e.g. CompressedGather)
+                const auto& restChilds = multiplyConvertNode->getChildEdges();
+                if (restChilds.empty())
+                    graph.RemoveEdge(multiplyConvertNode->getParentEdgeAt(0));
+            }
 
             graph.DropNode(convertNode);
             if (withSubtract)