[GPU] Moved RMSFusion higher in pipeline and added output type fuse

Author: Lyamin-Roman

src/common/transformations/include/ov_ops/rms.hpp

@@ -43,6 +43,10 @@ class TRANSFORMATIONS_API RMS : public ov::op::Op {
         m_epsilon = epsilon;
     }
 
+    void set_rms_output_type(const element::Type& output_type) {
+        m_output_type = output_type;
+    }
+
 private:
     double m_epsilon{0};
     ov::element::Type m_output_type;

src/common/transformations/src/transformations/common_optimizations/rms_fusion.cpp

@@ -13,6 +13,7 @@
 #include "openvino/op/reduce_mean.hpp"
 #include "openvino/op/sqrt.hpp"
 #include "openvino/pass/manager.hpp"
+#include "openvino/pass/pattern/op/or.hpp"
 #include "openvino/pass/pattern/op/wrap_type.hpp"
 #include "ov_ops/rms.hpp"
 #include "transformations/utils/utils.hpp"
@@ -57,11 +58,15 @@ RMSFusion::RMSFusion(bool force_tail_convert) {
     auto sqrt = wrap_type<ov::op::v0::Sqrt>({add_eps});
 
     // 1/Sqrt(ReduceMean(x^2,axes)+eps)
-    auto const_div = wrap_type<ov::op::v0::Constant>(constant_value(-1));
-    auto div = wrap_type<ov::op::v1::Power>({sqrt, const_div});
+    auto const_pow = wrap_type<ov::op::v0::Constant>(constant_value(-1));
+    auto pow = wrap_type<ov::op::v1::Power>({sqrt, const_pow});
+
+    auto const_div = wrap_type<ov::op::v0::Constant>(constant_value(1));
+    auto div = wrap_type<ov::op::v1::Divide>({const_div, sqrt});
+    auto div_or_pow = std::make_shared<pattern::op::Or>(OutputVector{div, pow});
 
     // x * 1/Sqrt(ReduceMean(x^2,axes)+eps)
-    auto mul1 = wrap_type<ov::op::v1::Multiply>({x, div});
+    auto mul1 = wrap_type<ov::op::v1::Multiply>({x, div_or_pow});
 
     // x * 1/Sqrt(ReduceMean(x^2,axes)+eps) * gamma
     auto gamma = wrap_type<ov::op::v0::Constant>(type_matches(element::f32));

src/common/transformations/src/transformations/convert_precision.cpp

@@ -12,6 +12,7 @@
 #include "openvino/pass/constant_folding.hpp"
 #include "openvino/pass/manager.hpp"
 #include "openvino/reference/convert.hpp"
+#include "ov_ops/rms.hpp"
 #include "ov_ops/type_relaxed.hpp"
 #include "transformations/fp16_compression/align_mixed_fp32_fp16_types.hpp"
 #include "transformations/fp16_compression/mark_decompression_convert_constant_folding.hpp"
@@ -59,6 +60,7 @@ bool fuse_type_to_maxpool(const std::shared_ptr<ov::Node>& node, const precision
 bool fuse_type_to_nonzero(const std::shared_ptr<ov::Node>& node, const precisions_map& precisions);
 bool fuse_type_to_bucketize(const std::shared_ptr<ov::Node>& node, const precisions_map& precisions);
 bool fuse_type_to_ctc_greedy_decoder_seq_len(const std::shared_ptr<ov::Node>& node, const precisions_map& precisions);
+bool fuse_type_to_rms(const std::shared_ptr<ov::Node>& node, const precisions_map& precisions);
 
 bool fuse_type_to_random_uniform_v8(const std::shared_ptr<ov::Node>& node, const precisions_map& precisions);
 
@@ -465,7 +467,8 @@ bool ov::pass::ConvertPrecision::run_on_model(const std::shared_ptr<ov::Model>&
         {ov::op::v0::PriorBox::get_type_info_static(), fuse_type_to_prior_box<ov::op::v0::PriorBox>},
         {ov::op::v8::PriorBox::get_type_info_static(), fuse_type_to_prior_box<ov::op::v8::PriorBox>},
         {ov::op::v0::PriorBoxClustered::get_type_info_static(), fuse_type_to_prior_box<ov::op::v0::PriorBoxClustered>},
-        {ov::op::v15::SearchSorted::get_type_info_static(), fuse_type_to_search_sorted_v15}};
+        {ov::op::v15::SearchSorted::get_type_info_static(), fuse_type_to_search_sorted_v15},
+        {ov::op::internal::RMS::get_type_info_static(), fuse_type_to_rms}};
 
     for (const auto& it : m_additional_type_to_fuse_map) {
         type_to_fuse[it.first] = it.second;
@@ -858,6 +861,20 @@ bool fuse_type_to_nms_rotated(const std::shared_ptr<ov::Node>& node, const preci
     return res;
 }
 
+bool fuse_type_to_rms(const std::shared_ptr<ov::Node>& node, const precisions_map& precisions) {
+    auto it = precisions.find(node->get_output_element_type(0));
+    if (it == precisions.end())
+        return false;
+    const auto& to = it->second;
+    if (auto rms = ov::as_type_ptr<ov::op::internal::RMS>(node)) {
+        if (to.is_real()) {
+            rms->set_rms_output_type(to);
+            return true;
+        }
+    }
+    return false;
+}
+
 namespace {
 
 bool update_type(size_t idx,

src/common/transformations/tests/common_optimizations/rms_norm_decomposition_test.cpp

@@ -77,6 +77,19 @@ TEST_F(TransformationTestsF, RMSNormFusionTest2) {
         model = std::make_shared<ov::Model>(ov::NodeVector{comp}, ov::ParameterVector{input});
         manager.register_pass<RMSFusion>();
     }
+    {
+        auto input = std::make_shared<ov::opset10::Parameter>(ov::element::f32, ov::Shape{1, 2, 6});
+
+        auto rms_const = ov::opset10::Constant::create(ov::element::f32,
+                                                       ov::Shape{6},
+                                                       {0.029f, 0.014f, 0.003f, 0.013f, 0.015f, 0.009f});
+        auto rms = std::make_shared<ov::op::internal::RMS>(input, rms_const, 1e-5f, ov::element::f16);
+
+        model_ref = std::make_shared<ov::Model>(ov::NodeVector{rms}, ov::ParameterVector{input});
+    }
+    comparator.enable(FunctionsComparator::CmpValues::ACCURACY);
+    comparator.enable(FunctionsComparator::CmpValues::CONST_VALUES);
+    comparator.enable(FunctionsComparator::CmpValues::ATTRIBUTES);
 }
 
 TEST_F(TransformationTestsF, RMSNormFusionTest3) {
@@ -113,7 +126,7 @@ TEST_F(TransformationTestsF, RMSNormFusionTest4) {
         auto eps = ov::opset10::Constant::create(ov::element::f32, {}, {1e-5f});
         auto add_eps = std::make_shared<ov::opset10::Add>(mean, eps);
         auto sqrt = std::make_shared<ov::opset10::Sqrt>(add_eps);
-        auto div_const = ov::opset10::Constant::create(ov::element::f32, {}, {1});
+        auto div_const = ov::opset10::Constant::create(ov::element::f32, {}, {-1});
         auto div = std::make_shared<ov::opset10::Divide>(div_const, sqrt);
         auto mul1 = std::make_shared<ov::opset10::Multiply>(input, div);
         auto gamma = ov::opset10::Constant::create(ov::element::f32,

src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp

@@ -113,6 +113,7 @@
 #include "transformations/op_conversions/convert_broadcast3.hpp"
 #include "transformations/op_conversions/convert_deformable_conv_v8_to_v1.hpp"
 #include "transformations/op_conversions/convert_depth_to_space.hpp"
+#include "transformations/op_conversions/convert_divide.hpp"
 #include "transformations/op_conversions/convert_gather_0d.hpp"
 #include "transformations/op_conversions/convert_gather_downgrade.hpp"
 #include "transformations/op_conversions/convert_gelu.hpp"
@@ -377,6 +378,16 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
             return !is_decompression_multiply(node, device_info.supports_immad);
         });
 
+        pass_config->set_callback<ov::pass::RMSFusion>([=](const_node_ptr& root) -> bool {
+            if (!root->get_input_partial_shape(0).is_static()) {
+                return false;
+            }
+            const auto& gamma_shape = root->get_input_partial_shape(0).to_shape();
+            const int32_t vec_size = 8;
+            return static_cast<int32_t>((gamma_shape.back() / vec_size)) > static_cast<int32_t>(device_info.max_work_group_size);
+        });
+        manager.register_pass<ov::pass::RMSFusion>(false);
+
         const bool keep_precision_sensitive_in_fp32_1 = true;
         const bool convert_input_output_precision = false;
         const bool store_original_precision_as_rt_attribute = true;
@@ -922,16 +933,6 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
 
         manager.register_pass<ov::pass::ConvertGatherToGatherCompressed>();
         auto pass_config = manager.get_pass_config();
-        pass_config->set_callback<ov::pass::RMSFusion>([=](const_node_ptr& root) -> bool {
-            if (!root->get_input_node_ptr(0)->get_input_partial_shape(0).is_static()) {
-                return false;
-            }
-            const auto& gamma_shape = root->get_input_node_ptr(0)->get_input_partial_shape(0).to_shape();
-            const int32_t vec_size = 8;
-            return static_cast<int32_t>((gamma_shape.back() / vec_size)) > static_cast<int32_t>(device_info.max_work_group_size);
-        });
-
-        manager.register_pass<ov::pass::RMSFusion>();
         manager.register_pass<ov::intel_gpu::KVCacheFusion>();
         manager.register_pass<ov::intel_gpu::FullyConnectedConvertFusion>();
         manager.register_pass<ov::intel_gpu::TransposeFusion>(device_info.supports_immad);
@@ -997,7 +998,6 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
         GPU_DEBUG_IF(cldnn::debug_configuration::get_instance()->verbose >= 1) {
             manager.register_pass<ov::intel_gpu::PrintModelStatistics>();
         }
-
         manager.run_passes(func);
     }
 }