[xpu] add c++ test of fused_multi_transformer_xpu_quant_pass

mayang002 · mayang002 · commit dde77c232cf7 · 2023-03-17T16:03:07.000+08:00
diff --git a/paddle/fluid/framework/ir/CMakeLists.txt b/paddle/fluid/framework/ir/CMakeLists.txt
@@ -235,7 +235,7 @@ if(WITH_XPU)
   pass_library(link_xpu_op_max_pass inference DIR xpu DEPS ${XPU_PASS_DEPS})
   pass_library(delete_isolated_node_pass inference DIR xpu DEPS
                ${XPU_PASS_DEPS})
-  pass_library(fused_multi_transformer_quant_pass inference DIR xpu DEPS
+  pass_library(fused_multi_transformer_xpu_quant_pass inference DIR xpu DEPS
                ${XPU_PASS_DEPS})
 endif()
 
@@ -495,4 +495,8 @@ if(WITH_XPU)
     test_delete_isolated_node_pass
     SRCS xpu/delete_isolated_node_pass_test.cc
     DEPS delete_isolated_node_pass)
+  cc_test(
+    test_fused_multi_transformer_xpu_quant_pass
+    SRCS xpu/fused_multi_transformer_xpu_quant_pass_tester.cc
+    DEPS fused_multi_transformer_xpu_quant_pass)
 endif()
diff --git a/paddle/fluid/framework/ir/fuse_multi_transformer_layer_pass_tester.cc b/paddle/fluid/framework/ir/fuse_multi_transformer_layer_pass_tester.cc
@@ -75,7 +75,7 @@ TEST(FuseMultiTransformerLayerPass, encoder_fp) {
         1,
         {2, -1, 16, 1024, 64},
         0);
-    auto* out = layers.fused_multi_transformer(x,
+    auto outs = layers.fused_multi_transformer(x,
                                                cache_kv,
                                                src_mask,
                                                qkv_w,
@@ -93,7 +93,7 @@ TEST(FuseMultiTransformerLayerPass, encoder_fp) {
                                                0.1,
                                                1e-12);
 
-    x = out;
+    x = outs[0];
   }
   std::unique_ptr<ir::Graph> graph(new ir::Graph(layers.main_program()));
   graph->Set("__param_scope__", CreateParamScope());
@@ -126,7 +126,7 @@ TEST(FuseMultiTransformerLayerPass, decoder_fp) {
   for (int i = 0; i < num_layers; ++i) {
     auto* shape_out = layers.shape(src_mask);
     auto* time_stamp = layers.slice(shape_out, {0}, {3}, {4});
-    auto* out = layers.fused_multi_transformer(x,
+    auto outs = layers.fused_multi_transformer(x,
                                                cache_kv,
                                                src_mask,
                                                qkv_w,
@@ -145,7 +145,7 @@ TEST(FuseMultiTransformerLayerPass, decoder_fp) {
                                                1e-12,
                                                time_stamp);
 
-    x = out;
+    x = outs[0];
   }
   std::unique_ptr<ir::Graph> graph(new ir::Graph(layers.main_program()));
   auto param_scope = CreateParamScope();
diff --git a/paddle/fluid/framework/ir/pass.cc b/paddle/fluid/framework/ir/pass.cc
@@ -49,7 +49,7 @@ static const std::vector<std::string> support_subgraph_passes = {
     "fuse_multi_transformer_layer_pass",
     "delete_quant_dequant_linear_op_pass",
     "delete_weight_dequant_linear_op_pass",
-    "fused_multi_transformer_quant_pass",
+    "fused_multi_transformer_xpu_quant_pass",
     "fc_xpu_fuse_pass",
     "delete_op_device_pass"};
 
diff --git a/paddle/fluid/framework/ir/pass_tester_helper.h b/paddle/fluid/framework/ir/pass_tester_helper.h
@@ -571,33 +571,35 @@ struct Layers {
     return out;
   }
 
-  VarDesc* fused_multi_transformer(VarDesc* x,
-                                   VarDesc* cache_kv,
-                                   VarDesc* src_mask,
-                                   VarDesc* qkv_w,
-                                   VarDesc* qkv_bias,
-                                   VarDesc* out_linear_w,
-                                   VarDesc* out_linear_bias,
-                                   VarDesc* ffn1_w,
-                                   VarDesc* ffn1_bias,
-                                   VarDesc* ffn2_w,
-                                   VarDesc* ffn2_bias,
-                                   VarDesc* ln_scale,
-                                   VarDesc* ln_bias,
-                                   VarDesc* ffn_ln_scale,
-                                   VarDesc* ffn_ln_bias,
-                                   float epsilon,
-                                   float dropout_rate,
-                                   VarDesc* time_stamp = nullptr,
-                                   VarDesc* qkv_out_scale = nullptr,
-                                   VarDesc* out_linear_out_scale = nullptr,
-                                   VarDesc* ffn1_out_scale = nullptr,
-                                   VarDesc* ffn2_out_scale = nullptr,
-                                   std::vector<float> qkv_in_scale = {},
-                                   std::vector<float> out_linear_in_scale = {},
-                                   std::vector<float> ffn1_in_scale = {},
-                                   std::vector<float> ffn2_in_scale = {}) {
+  std::vector<VarDesc*> fused_multi_transformer(
+      VarDesc* x,
+      VarDesc* cache_kv,
+      VarDesc* src_mask,
+      VarDesc* qkv_w,
+      VarDesc* qkv_bias,
+      VarDesc* out_linear_w,
+      VarDesc* out_linear_bias,
+      VarDesc* ffn1_w,
+      VarDesc* ffn1_bias,
+      VarDesc* ffn2_w,
+      VarDesc* ffn2_bias,
+      VarDesc* ln_scale,
+      VarDesc* ln_bias,
+      VarDesc* ffn_ln_scale,
+      VarDesc* ffn_ln_bias,
+      float epsilon,
+      float dropout_rate,
+      VarDesc* time_stamp = nullptr,
+      VarDesc* qkv_out_scale = nullptr,
+      VarDesc* out_linear_out_scale = nullptr,
+      VarDesc* ffn1_out_scale = nullptr,
+      VarDesc* ffn2_out_scale = nullptr,
+      std::vector<float> qkv_in_scale = {},
+      std::vector<float> out_linear_in_scale = {},
+      std::vector<float> ffn1_in_scale = {},
+      std::vector<float> ffn2_in_scale = {}) {
     VarDesc* out = lod_tensor(unique_name());
+    VarDesc* cache_kv_out = lod_tensor(unique_name());
     OpDesc* op = program_.MutableBlock(0)->AppendOp();
     std::string op_type = qkv_out_scale ? "fused_multi_transformer_int8"
                                         : "fused_multi_transformer";
@@ -623,6 +625,7 @@ struct Layers {
     op->SetAttr("dropout_rate", dropout_rate);
     op->SetAttr("epsilon", epsilon);
     op->SetOutput("Out", {out->Name()});
+    op->SetOutput("CacheKVOut", {cache_kv_out->Name()});
 
     if (time_stamp) {
       op->SetInput("TimeStep", {time_stamp->Name()});
@@ -638,7 +641,8 @@ struct Layers {
       op->SetAttr("ffn1_in_scale", ffn1_in_scale);
       op->SetAttr("ffn2_in_scale", ffn2_in_scale);
     }
-    return out;
+    std::vector<VarDesc*> outs = {out, cache_kv_out};
+    return outs;
   }
 
   VarDesc* dequantize_linear(VarDesc* x,
diff --git a/paddle/fluid/framework/ir/xpu/fused_multi_transformer_xpu_quant_pass.cc b/paddle/fluid/framework/ir/xpu/fused_multi_transformer_xpu_quant_pass.cc
@@ -250,7 +250,7 @@ FusedMultiTransformerPattern::FusedMultiTransformerPattern(
 1. transpose and quantify the weights of fused_multi_transformer op from fp32 to
 int16
 */
-class FusedMultiTransformerQuantPass : public FusePassBase {
+class FusedMultiTransformerXPUQuantPass : public FusePassBase {
  protected:
   void ApplyImpl(ir::Graph* graph) const override;
 
@@ -263,32 +263,30 @@ class FusedMultiTransformerQuantPass : public FusePassBase {
                 bool with_seq_lengths,
                 bool with_src_mask) const;
 
-  const std::string name_scope_{"fused_multi_transformer_quant_pass"};
+  const std::string name_scope_{"fused_multi_transformer_xpu_quant_pass"};
 };
 
-void FusedMultiTransformerQuantPass::ApplyImpl(ir::Graph* graph) const {
+void FusedMultiTransformerXPUQuantPass::ApplyImpl(ir::Graph* graph) const {
   PADDLE_ENFORCE_NOT_NULL(
       graph, platform::errors::PreconditionNotMet("graph should not be null."));
   Init(name_scope_, graph);
-  VLOG(3) << "DEBUG: in FusedMultiTransformerQuantPass::ApplyImpl";
+  VLOG(3) << "in FusedMultiTransformerXPUQuantPass::ApplyImpl";
 
   int found_subgraph_count = 0;
-  for (bool with_cache_kv : {true, false}) {
-    for (bool with_time_step : {true, false}) {
-      found_subgraph_count += ApplyImpl(
-          graph, with_cache_kv, false, false, with_time_step, false, true);
-    }
+  for (bool with_time_step : {true, false}) {
+    found_subgraph_count +=
+        ApplyImpl(graph, true, false, false, with_time_step, false, true);
   }
   AddStatis(found_subgraph_count);
 }
 
-int FusedMultiTransformerQuantPass::ApplyImpl(ir::Graph* graph,
-                                              bool with_cache_kv,
-                                              bool with_pre_caches,
-                                              bool with_rotary_pos_emb,
-                                              bool with_time_step,
-                                              bool with_seq_lengths,
-                                              bool with_src_mask) const {
+int FusedMultiTransformerXPUQuantPass::ApplyImpl(ir::Graph* graph,
+                                                 bool with_cache_kv,
+                                                 bool with_pre_caches,
+                                                 bool with_rotary_pos_emb,
+                                                 bool with_time_step,
+                                                 bool with_seq_lengths,
+                                                 bool with_src_mask) const {
   GraphPatternDetector gpd;
   patterns::FusedMultiTransformerPattern pattern(gpd.mutable_pattern(),
                                                  name_scope_,
@@ -302,7 +300,7 @@ int FusedMultiTransformerQuantPass::ApplyImpl(ir::Graph* graph,
   int found_subgraph_count = 0;
   auto handler = [&](const GraphPatternDetector::subgraph_t& subgraph,
                      Graph* graph) {
-    VLOG(4) << "handle FusedMultiTransformerQuantPass fuse";
+    VLOG(4) << "handle FusedMultiTransformerXPUQuantPass fuse";
 
     GET_IR_NODE(x);
     GET_IR_NODE(ln_scale);
@@ -544,5 +542,5 @@ int FusedMultiTransformerQuantPass::ApplyImpl(ir::Graph* graph,
 }  // namespace framework
 }  // namespace paddle
 
-REGISTER_PASS(fused_multi_transformer_quant_pass,
-              paddle::framework::ir::FusedMultiTransformerQuantPass);
+REGISTER_PASS(fused_multi_transformer_xpu_quant_pass,
+              paddle::framework::ir::FusedMultiTransformerXPUQuantPass);
diff --git a/paddle/fluid/framework/ir/xpu/fused_multi_transformer_xpu_quant_pass_tester.cc b/paddle/fluid/framework/ir/xpu/fused_multi_transformer_xpu_quant_pass_tester.cc
@@ -0,0 +1,170 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <gtest/gtest.h>
+
+#include "paddle/fluid/framework/ir/pass.h"
+#include "paddle/fluid/framework/ir/pass_tester_helper.h"
+
+#define DEF_INPUT_DATA                                                  \
+  Layers layers;                                                        \
+  auto* x = layers.data("x", {1, 128, 1024});                           \
+  auto* src_mask = layers.data("src_mask", {1, 16, 128, 128});          \
+  auto* ln_scale = layers.data("ln_scale", {1024}, true);               \
+  auto* ln_bias = layers.data("ln_bias", {1024}, true);                 \
+  auto* qkv_w = layers.data("qkv_w", {3, 16, 64, 1024}, true);          \
+  auto* qkv_bias = layers.data("qkv_bias", {3, 16, 64}, true);          \
+  auto* out_linear_w = layers.data("out_linear_w", {1024, 1024}, true); \
+  auto* out_linear_bias = layers.data("out_linear_bias", {1024}, true); \
+  auto* ffn_ln_scale = layers.data("ffn_ln_scale", {1024}, true);       \
+  auto* ffn_ln_bias = layers.data("ffn_ln_bias", {1024}, true);         \
+  auto* ffn1_w = layers.data("ffn1_w", {1024, 4096}, true);             \
+  auto* ffn1_bias = layers.data("ffn1_bias", {4096}, true);             \
+  auto* ffn2_w = layers.data("ffn2_w", {4096, 1024}, true);             \
+  auto* ffn2_bias = layers.data("ffn2_bias", {1024}, true);
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+void AddVarToScope(Scope* param_scope,
+                   const std::string& name,
+                   const DDim& dims) {
+  auto* tensor = param_scope->Var(name)->GetMutable<phi::DenseTensor>();
+  tensor->Resize(dims);
+  tensor->mutable_data<float>(platform::CPUPlace());
+}
+
+Scope* CreateParamScope() {
+  auto param_scope = new Scope();
+  AddVarToScope(param_scope, "ln_scale", {1024});
+  AddVarToScope(param_scope, "ln_bias", {1024});
+  AddVarToScope(param_scope, "ffn_ln_scale", {1024});
+  AddVarToScope(param_scope, "ffn_ln_bias", {1024});
+
+  AddVarToScope(param_scope, "qkv_w", {3, 16, 64, 1024});
+  AddVarToScope(param_scope, "out_linear_w", {1024, 1024});
+  AddVarToScope(param_scope, "ffn1_w", {1024, 4096});
+  AddVarToScope(param_scope, "ffn2_w", {4096, 1024});
+  AddVarToScope(param_scope, "qkv_bias", {3072});
+  AddVarToScope(param_scope, "out_linear_bias", {1024});
+  AddVarToScope(param_scope, "ffn1_bias", {4096});
+  AddVarToScope(param_scope, "ffn2_bias", {1024});
+
+  return param_scope;
+}
+
+TEST(FusedMultiTransformerXPUQuantPass, context_stage) {
+  DEF_INPUT_DATA
+
+  auto* cache_kv = layers.fill_constant_batch_size_like(
+      x,
+      static_cast<int>(proto::VarType::FP32),
+      0,
+      1,
+      {2, -1, 16, 1024, 64},
+      0);
+
+  layers.fused_multi_transformer(x,
+                                 cache_kv,
+                                 src_mask,
+                                 qkv_w,
+                                 qkv_bias,
+                                 out_linear_w,
+                                 out_linear_bias,
+                                 ffn1_w,
+                                 ffn1_bias,
+                                 ffn2_w,
+                                 ffn2_bias,
+                                 ln_scale,
+                                 ln_bias,
+                                 ffn_ln_scale,
+                                 ffn_ln_bias,
+                                 0.1,
+                                 1e-12);
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(layers.main_program()));
+  graph->Set("__param_scope__", CreateParamScope());
+
+  auto pass =
+      PassRegistry::Instance().Get("fused_multi_transformer_xpu_quant_pass");
+  if (pass.get() == nullptr) {
+    LOG(INFO) << "get fused_multi_transformer_xpu_quant_pass failed";
+  }
+
+  graph.reset(pass->Apply(graph.release()));
+  int num_nodes_after = GetNumOpNodes(graph, "fused_multi_transformer_xpu");
+  VLOG(3) << DebugString(graph);
+
+  PADDLE_ENFORCE_EQ(
+      num_nodes_after,
+      1,
+      platform::errors::InvalidArgument(
+          "After the fuse_multi_transformer_layer_pass, "
+          "The node num in graph should be 1, but the result is %d",
+          num_nodes_after));
+}
+
+TEST(FusedMultiTransformerXPUQuantPass, decoder_stage) {
+  DEF_INPUT_DATA
+
+  auto* cache_kv = layers.fill_constant_batch_size_like(
+      x,
+      static_cast<int>(proto::VarType::FP32),
+      0,
+      1,
+      {2, -1, 16, 1024, 64},
+      0);
+  auto* time_step = layers.data("time_step", {1});
+  layers.fused_multi_transformer(x,
+                                 cache_kv,
+                                 src_mask,
+                                 qkv_w,
+                                 qkv_bias,
+                                 out_linear_w,
+                                 out_linear_bias,
+                                 ffn1_w,
+                                 ffn1_bias,
+                                 ffn2_w,
+                                 ffn2_bias,
+                                 ln_scale,
+                                 ln_bias,
+                                 ffn_ln_scale,
+                                 ffn_ln_bias,
+                                 0.1,
+                                 1e-12,
+                                 time_step);
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(layers.main_program()));
+  graph->Set("__param_scope__", CreateParamScope());
+
+  auto pass =
+      PassRegistry::Instance().Get("fused_multi_transformer_xpu_quant_pass");
+  if (pass.get() == nullptr) {
+    LOG(INFO) << "get fused_multi_transformer_xpu_quant_pass failed";
+  }
+
+  graph.reset(pass->Apply(graph.release()));
+  int num_nodes_after = GetNumOpNodes(graph, "fused_multi_transformer_xpu");
+  VLOG(3) << DebugString(graph);
+
+  PADDLE_ENFORCE_EQ(
+      num_nodes_after,
+      1,
+      platform::errors::InvalidArgument(
+          "After the fuse_multi_transformer_layer_pass, "
+          "The node num in graph should be 1, but the result is %d",
+          num_nodes_after));
+}
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+USE_PASS(fused_multi_transformer_xpu_quant_pass);
diff --git a/paddle/fluid/inference/api/paddle_pass_builder.cc b/paddle/fluid/inference/api/paddle_pass_builder.cc
@@ -524,7 +524,7 @@ XpuPassStrategy::XpuPassStrategy() : PassStrategy({}) {
       "embedding_with_eltwise_add_xpu_fuse_pass",
       "multi_encoder_xpu_fuse_pass",
       "multi_encoder_xpu_slice_fuse_pass",
-      "fused_multi_transformer_quant_pass",
+      "fused_multi_transformer_xpu_quant_pass",
       "fc_xpu_fuse_pass",
       "link_xpu_op_max_pass",
       "delete_op_device_pass",
