Fused attention op forward #35905

limin2021 · 2021-09-22T07:12:13Z

PR types

New features

PR changes

OPs

Describe

功能：本PR的目标是提高attention模块的计算性能。
为了减少框架层对op的调度开销，本PR通过在C++层手动实现attention模块，对外提供attention 大op；
为了减少防存开销，本PR采取了两种优化方法：
（1）在q,k,v计算时通过共享输入X，将该处的gemm，transpose和bias add从三次调用减少为一次；
（2）使用kernel融合优化技术，在不同cuda kernel之间通过寄存器传输数据；
fused_attention_op 实现的计算逻辑：
fused_attention_op与paddle已有的MultiHeadAttention layer的不同：
（1）计算逻辑范围扩大了，详见上面的伪代码。
（2）q, k, v的weight存储格式不一样。
原有的：保存在三个weight张量中，WQ, WK, WV
本PR：保存在一个weight张量中，qkv_weight
由WQ, WK, WV得到qkv_weight的方法：
实现:
本PR是fused_attention_op 的前向实现，具体细节：

(1) fused_attention_op.cc and fused_attention_op.cu
The C++ forward impl for fused_attention_op. The impl uses these PRs:
#34883, #35308, #35350 #35621 , #35903, #36185

(2) functional/fused_transformer.py
The python api for fused_attention_op.
Here, it only include dynamic graph api,
the static graph api will be added in the next PR.

(3) test_fused_attention_op.py
The unittest script for fused_attention_op: dynamic, forward;

(4) paddle/fluid/operators/dropout_impl_util.h
Modifications of contents of dropout_impl_util.h in #35820 is overlapped by #36185.
In this PR, we recovered the contents to be same as #35820.

(5) Fix bugs: remove useless "print" in framework.py.

Unittest:

paddle-bot-old · 2021-09-22T07:12:17Z

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

zkh2016 · 2021-09-23T08:51:35Z

paddle/fluid/operators/fused/fused_attention_op.cc

+             "H(`begin_norm_axis` splits the tensor(`X`) to a matrix [N,H])."
+             "It is applied to the output.")
+        .AsDispensable();
+    // AddInput("Seed",


注释代码可以删除

zkh2016 · 2021-09-23T08:52:23Z

paddle/fluid/operators/fused/fused_attention_op.cc

+        });
+
+    AddComment(R"DOC(
+Fused attention op: 


格式似乎不对

zkh2016 · 2021-09-23T08:55:37Z

paddle/fluid/operators/fused/fused_attention_op.h

+
+namespace paddle {
+namespace operators {}  // namespace operators
+}  // namespace paddle


这个文件可以不需要吧

是需要的。pre-commit要求op.cc和op.cu必须include op.h，否则过不了。

应该不是强制要求必须有.h的，multi_dot就只有multi_dot_op.cc。

zkh2016 · 2021-09-23T08:56:25Z

python/paddle/fluid/tests/unittests/test_fused_attention_op.py

+import unittest
+
+
+def _convert_attention_mask(attn_mask, dtype):


这个可以调transformer里面的

zhangting2020 · 2021-09-23T09:13:31Z

paddle/fluid/operators/fused/fused_attention_op.cc

+                          "(batch_size, seq_len, dim_embed),"
+                          "but received dimensions of"
+                          "Input is [%d]",
+                          x_dim.size()));


这里的QKV_input指的是是x？最好对应到用户接口里的输入。报错信息需要让用户看到后能准确理解和定位问题。
下面也有类似问题

Done. 已全部check。

zhangting2020 · 2021-09-23T09:18:53Z

paddle/fluid/operators/fused/fused_dropout_helper.h

@@ -0,0 +1,298 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.


这个是和 #35843 完全一样的吗？

一样的哈。等他的合入，这里merge即可。此处不需要review。

TCChenlong

LGTM

zkh2016

LGTM

xingfeng01 · 2021-09-24T02:29:58Z

LGTM

zhiqiu

LGTM for op_function_generator.cc

Xreki · 2021-09-24T05:44:08Z

paddle/fluid/operators/fused/fused_attention_op.cc

+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


注意copyright的格式

Xreki · 2021-09-24T05:45:32Z

paddle/fluid/operators/fused/fused_attention_op.cc

+    OP_INOUT_CHECK(ctx->HasInput("QKVW"), "Input", "QKVW", "FusedAttentionOp");
+    OP_INOUT_CHECK(ctx->HasInput("QKVBias"), "Input", "QKVBias",
+                   "FusedAttentionOp");
+    OP_INOUT_CHECK(ctx->HasInput("OutLinearW"), "Input", "OutLinearW",


这是一个输入，名字叫OutXxx会不会有点奇怪？

在下一个PR： fused_attention_bw中一起修改，会将OutLinear修改成Linear，也会一起polish其他命名。

Xreki · 2021-09-24T05:52:32Z

paddle/fluid/operators/fused/fused_attention_op.cc

+              platform::errors::InvalidArgument(
+                  "'attn_dropout_prob' must be between 0.0 and 1.0."));
+        });
+    AddAttr<bool>("is_test1",


尽量不要叫xxx_1这种名字，很难区分。

Xreki · 2021-09-24T05:54:06Z

paddle/fluid/operators/fused/fused_attention_op.cc

+        "   train: out = input * mask / ( 1.0 - dropout_prob )"
+        "   inference: out = input"
+        "   dropout op can be removed from the program. the program will be "
+        "efficient")


这么长一段重复的解释，用一个宏、或者一个函数？

Xreki · 2021-09-24T05:55:50Z

paddle/fluid/operators/fused/fused_attention_op.cu

+#ifdef __HIPCC__
+#include <hipcub/hipcub.hpp>
+namespace cub = hipcub;
+#endif


cmake里面加了限制，不支持ROCM。

Xreki · 2021-09-24T06:05:19Z

paddle/fluid/operators/fused/fused_dropout_helper.h

+    LaunchResidualDropoutBiasGrad<T, uint8_t>(
+        dout, mask, dropout_param_.dropout_prob,
+        dropout_param_.is_upscale_in_train, rows_, cols_, dsrc, dbias, ctx);
+    cudaMemcpyAsync(dresidual, dout, rows_ * cols_ * sizeof(T),


用memory::Copy

这个文件引自#35843 ，可在其下review，此处无需review。待#35843 合入后，本PR会merge更新。

Xreki · 2021-09-24T06:06:27Z

python/paddle/fluid/tests/unittests/test_fused_attention_op.py

+
+@unittest.skipIf(not core.is_compiled_with_cuda(),
+                 "Paddle core is not compiled with CUDA")
+class TestFusedAttentionOp(unittest.TestCase):


OP单测建议继承OpTest。

Xreki · 2021-09-24T06:07:21Z

python/paddle/fluid/tests/unittests/test_fused_attention_op.py

+import unittest
+
+
+@unittest.skipIf(not core.is_compiled_with_cuda(),


只支持CUDA，可以加在CMakeLists.txt里面。

Xreki · 2021-09-24T06:09:43Z

python/paddle/nn/functional/__init__.py

@@ -60,6 +60,7 @@
 from .conv import conv1d  # noqa: F401
 from .conv import conv1d_transpose  # noqa: F401
 from .common import linear  # noqa: F401
+from .common import fused_multihead_attention  # noqa: F401


这个functional接口是否要暴露给用户？

需要暴露，已move到functional/fused_transformer.py中。

Xreki · 2021-09-24T06:10:40Z

python/paddle/nn/functional/common.py

@@ -1502,6 +1502,33 @@ def linear(x, weight, bias=None, name=None):
        return res


+def fused_multihead_attention(x,


建议还是先实现在fused_transformer.py里面吧。

zkh2018 · 2021-09-26T06:33:20Z

paddle/fluid/operators/fused/fused_attention_op.cu

+    auto *qkv_bias_out_data = qkv_bias_out->mutable_data<T>(ctx.GetPlace());
+
+    // get data ptr for FMHA.
+    auto *src_mask_data = (src_mask == nullptr ? nullptr : src_mask->data<T>());


这个变量好像没有使用的地方？

… fused_attention_op_2_fw

zhangting2020

LGTM

zhangting2020 · 2021-10-21T10:57:52Z

python/paddle/nn/functional/fused_transformer.py

+    out = out * v;
+    out = transpose(out, perm=[0, 2, 1, 3]);      
+	out = out_linear(out);
+	out = layer_norm(x + dropout(linear_bias + out));


这里格式是不是有点问题

zhangting2020 · 2021-10-21T10:58:56Z

python/paddle/nn/functional/fused_transformer.py

+                None, None, None, None, 1e-5, qkv_bias,
+                linear_bias, attn_mask)
+            # [2, 4, 128]
+            print(output)


这里[2, 4, 128]是输出shape吧，print(output.shape)?

TCChenlong

TODO：Fix API Docs

zhiqiu

LGTM for op_function_generator.cc

功能：本PR的目标是提高attention模块的计算性能。为了减少框架层对op的调度开销，本PR通过在C++层手动实现attention模块，对外提供attention 大op；为了减少防存开销，本PR采取了两种优化方法：（1）在q,k,v计算时通过共享输入X，将该处的gemm，transpose和bias add从三次调用减少为一次；（2）使用kernel融合优化技术，在不同cuda kernel之间通过寄存器传输数据；

…ubate (#36704) 将 #35905 和 #35843 PR中新增的的python api接口移到incubate目录下。

…ubate (PaddlePaddle#36704) 将 PaddlePaddle#35905 和 PaddlePaddle#35843 PR中新增的的python api接口移到incubate目录下。

* add op: fused_feedforward(backward) (#35611) 这个PR是fused_feedforward反向的代码相关kernel实现：fused_dropout_act_bias, fused_residual_dropout_bias, fused_layernorm_residual_dropout_bias fused_feedforward是一个融合算子，该算子对transformer模型的feed forward层的算子进行融合和封装，使得前端只呈现一个接口，通过融合减少部分访存和kernel launch的时间，以此提升性能。 * Move fused_attention and fused_feedforward functional api path to incubate (#36704) 将 #35905 和 #35843 PR中新增的的python api接口移到incubate目录下。

limin2021 added 3 commits September 22, 2021 05:01

Add fused_attention_op: add impl wrappers.

f5eee9f

Add fused_attention_op: forward.

e16e3b3

Add fused_attention_op: forward impl.

42f0372

limin2021 and others added 8 commits September 22, 2021 08:07

Remove useless code.

c6aebef

Remove useless code.

2c0ab6c

Remove docs.

ece3c08

Minors.

b18b405

Minors.

b939159

Update test_fused_attention_op.py

07fd753

Merge branch 'PaddlePaddle:develop' into fused_attention_op_2_fw

ef89a94

Remove static construction of python api.

b44d882

zkh2016 reviewed Sep 23, 2021

View reviewed changes

zhangting2020 reviewed Sep 23, 2021

View reviewed changes

Modifications accordding to reviews.

ff3df46

TCChenlong previously approved these changes Sep 23, 2021

View reviewed changes

kolinwei previously approved these changes Sep 24, 2021

View reviewed changes

zkh2016 previously approved these changes Sep 24, 2021

View reviewed changes

xingfeng01 previously approved these changes Sep 24, 2021

View reviewed changes

zhangting2020 previously approved these changes Sep 24, 2021

View reviewed changes

limin2021 mentioned this pull request Sep 24, 2021

Fused attention op backward #35935

Closed

lanxianghit previously approved these changes Sep 24, 2021

View reviewed changes

zhiqiu previously approved these changes Sep 24, 2021

View reviewed changes

Xreki reviewed Sep 24, 2021

View reviewed changes

zkh2018 reviewed Sep 26, 2021

View reviewed changes

Modifications accordding to Xreki's review.

8a4c2a8

limin2021 dismissed stale reviews from zhiqiu, lanxianghit, and zhangting2020 via 8a4c2a8 September 26, 2021 14:19

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

ae875ca

… fused_attention_op_2_fw

zhangting2020 previously approved these changes Oct 14, 2021

View reviewed changes

xingfeng01 previously approved these changes Oct 14, 2021

View reviewed changes

limin2021 mentioned this pull request Oct 18, 2021

Add fused attention op backward and python layer. #36498

Merged

Add "#require gpu" for sample code in english doc.

10687a6

limin2021 dismissed stale reviews from xingfeng01 and zhangting2020 via 10687a6 October 21, 2021 06:02

xingfeng01 previously approved these changes Oct 21, 2021

View reviewed changes

zhangting2020 reviewed Oct 21, 2021

View reviewed changes

Improve format of sample code.

0f93775

limin2021 dismissed xingfeng01’s stale review via 0f93775 October 21, 2021 11:08

zhangting2020 approved these changes Oct 21, 2021

View reviewed changes

xingfeng01 approved these changes Oct 21, 2021

View reviewed changes

TCChenlong approved these changes Oct 22, 2021

View reviewed changes

zkh2016 approved these changes Oct 22, 2021

View reviewed changes

zhiqiu approved these changes Oct 22, 2021

View reviewed changes

lanxianghit approved these changes Oct 22, 2021

View reviewed changes

lanxianghit merged commit d490621 into PaddlePaddle:develop Oct 22, 2021

limin2021 mentioned this pull request Oct 25, 2021

Move fused_attention and fused_feedforward functional api path to incubate #36704

Merged

lanxianghit pushed a commit that referenced this pull request Oct 26, 2021

Move fused_attention and fused_feedforward functional api path to inc…

9aeca2f

…ubate (#36704) 将 #35905 和 #35843 PR中新增的的python api接口移到incubate目录下。

zkh2016 mentioned this pull request Oct 26, 2021

Fused transformer encoder layer and fused feedforward layer #36604

Merged

limin2021 mentioned this pull request Oct 26, 2021

[cherry-pick-2.2]Add fused attention op backward and python layer. #36752

Merged

		import unittest


		def _convert_attention_mask(attn_mask, dtype):

		@@ -0,0 +1,298 @@
		/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.

		import unittest


		@unittest.skipIf(not core.is_compiled_with_cuda(),

		@@ -1502,6 +1502,33 @@ def linear(x, weight, bias=None, name=None):
		return res


		def fused_multihead_attention(x,

Fused attention op forward #35905

Fused attention op forward #35905

Conversation

limin2021 commented Sep 22, 2021 • edited Loading

PR types

PR changes

Describe

paddle-bot-old bot commented Sep 22, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

TCChenlong left a comment

Choose a reason for hiding this comment

zkh2016 left a comment

Choose a reason for hiding this comment

xingfeng01 commented Sep 24, 2021

zhiqiu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

limin2021 Sep 27, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

limin2021 Sep 26, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zhangting2020 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

TCChenlong left a comment

Choose a reason for hiding this comment

zhiqiu left a comment

Choose a reason for hiding this comment

limin2021 commented Sep 22, 2021 •

edited

Loading

limin2021 Sep 27, 2021 •

edited

Loading

limin2021 Sep 26, 2021 •

edited

Loading