[hybrid performance] softmax mask fuse op

[FeixLiu]

PR types

New features

PR changes

OPs

Describe

fuse mask elementwise add and softmax together, for transformer used

general pass:

# declare Q, K, dk and mask
QK = matmul(Q, K.transpose()) / sqrt(dk)    # launch first op for matmul
tmp = QK + mask                             # launch second op for elementwise_mul
res = softmax(tmp)                          # launch third op for softmax

fused pass:

# declare Q, K, dk and mask
QK = matmul(Q, K.transpose()) / sqrt(dk)    # launch first op for matmul
res = softmax_mask_fuse(QK, mask)           # launch second op for softmax_mask_fuse

performance, based on PaddleNLP/GPT under AMP:

model size	hybird size (dp, pp, mp)	no fuse throughput	fused throughput	gain
117M	8(1,1,8)	12318.38	15046.72	+22%
117M	1(1,1,1)	8729.29	9859.29	+13%
345M	8(1,1,1)	5317.19	6766.67	+27%
345M	1(1,1,1)	3796.09	4322.57	+14%
1.3B	8(1,1,8)	3970.90	4727.14	+19%

loss curve, for PaddleNLP/GPT under AMP:

average loss diff for 20,000 steps: 0.0077
loss diff between fused pass and no fuse pass:

Currently, this OP only supports fp16 dtype.

To use this op from python side, follow these codes for static mode:

import paddle.incubate as incubate
import paddle.fluid as fluid
import numpy as np

input_x = fluid.data(name="x", shape=[4, 4, 8, 132], dtype="float16")
input_mask = fluid.data(name="mask", shape=[4, 1, 8, 132], dtype="float16")
rst = incubate.softmax_mask_fuse(input_x, input_mask)

x_in_np = np.random.random((4, 4, 8, 132)).astype("float16")
mask = np.random.randint(0, 2, (4, 1, 8, 132)).astype("float16")
mask_in_np = np.where(mask == 1, -10000.0, mask)

exe = fluid.Executor(place)
fetches = exe.run(fluid.default_main_program(),
                  feed={"x": x_in_np,
                        "mask": mask_in_np},
                  fetch_list=[rst])

Follow these codes for dynamic mode:

x_in_np = np.random.random((2, 8, 8, 2040)).astype("float16")
mask = np.random.randint(0, 2, (2, 1, 8, 2040)).astype("float16")
mask_in_np = np.where(mask == 1, -10000.0, mask)
input_x = fluid.dygraph.to_variable(x_in_np)
input_mask = fluid.dygraph.to_variable(mask_in_np)

rst = incubate.softmax_mask_fuse(input_x, input_mask)

[paddle-bot-old]

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

[ForFishes]

LGTM

[Xreki]

LGTM for the modification of atol and LGTM for op benchmark ci.

[paddle-bot-old]

Sorry to inform you that 170cde8's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.

[ForFishes]

LGTM

[XiaoguangHu01]

LG API

[PangHua]

LGTM