wehub-resource-sync
621 lines
20 KiB
Python
621 lines
20 KiB
Python
# Copyright (c) 2021 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.
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from op_test import get_device_place
|
|
|
|
import paddle
|
|
from paddle.incubate.nn.layer.fused_transformer import FusedMultiHeadAttention
|
|
|
|
|
|
def fc(x, weight):
|
|
return np.matmul(x, weight)
|
|
|
|
|
|
def softmax(x):
|
|
np.seterr(invalid='ignore')
|
|
output = np.zeros(x.shape, dtype=np.float64)
|
|
for i in range(x.shape[0]):
|
|
for j in range(x.shape[1]):
|
|
for k in range(x.shape[2]):
|
|
x_curr = x[i, j, k, :]
|
|
e_x = np.exp(x_curr - np.amax(x_curr))
|
|
output[i, j, k, :] = e_x / np.sum(e_x)
|
|
return output
|
|
|
|
|
|
def batch_matmul(x, y):
|
|
assert x.shape[0] == y.shape[0]
|
|
assert x.shape[1] == y.shape[1]
|
|
retval = np.zeros(
|
|
(x.shape[0], x.shape[1], x.shape[2], y.shape[3]), dtype=np.float64
|
|
)
|
|
for i in range(x.shape[0]):
|
|
for j in range(x.shape[1]):
|
|
retval[i, j, :, :] = np.matmul(x[i, j, :, :], y[i, j, :, :])
|
|
return retval
|
|
|
|
|
|
def layer_norm(x, has_scale, has_bias, weight, bias, epsilon=1e-05):
|
|
batch_size, src_len, d_model = x.shape
|
|
x = x.reshape((batch_size * src_len, d_model))
|
|
mu = np.mean(x, axis=1, keepdims=True)
|
|
sigma_square = np.sum(np.square(x - mu), axis=1) / d_model
|
|
x1_up = x - mu
|
|
x1_down_1 = sigma_square + epsilon
|
|
x1_down = np.sqrt(x1_down_1)
|
|
x1_down = x1_down.reshape((x1_down.shape[0], 1))
|
|
x1 = x1_up / x1_down
|
|
x_scaled = x1
|
|
if has_scale:
|
|
x_scaled = weight * x1
|
|
x_scaled_bias = x_scaled
|
|
if has_bias:
|
|
x_scaled_bias = x_scaled + bias
|
|
x_scaled_bias = x_scaled_bias.reshape((batch_size, src_len, d_model))
|
|
return x_scaled_bias
|
|
|
|
|
|
def compute_reference(
|
|
pre_layer_norm,
|
|
query,
|
|
attn_mask,
|
|
ln_scale,
|
|
ln_bias,
|
|
ln_2_scale,
|
|
ln_2_bias,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
out_linear_weight,
|
|
out_linear_bias,
|
|
num_head,
|
|
transpose_qkv_wb,
|
|
):
|
|
batch_size = query.shape[0]
|
|
seq_len = query.shape[1]
|
|
embed_dim = query.shape[2]
|
|
|
|
has_bias = True
|
|
if ln_bias is None:
|
|
has_bias = False
|
|
|
|
if pre_layer_norm:
|
|
ln_out = layer_norm(query, True, has_bias, ln_scale, ln_bias)
|
|
|
|
head_dim = embed_dim // num_head
|
|
if not transpose_qkv_wb:
|
|
# embed_dim, 3, num_heads, self.head_dim
|
|
qkv_weight = qkv_weight.transpose((3, 0, 1, 2))
|
|
qkv_weight = qkv_weight.reshape(
|
|
qkv_weight.shape[0],
|
|
qkv_weight.shape[1] * qkv_weight.shape[2] * qkv_weight.shape[3],
|
|
)
|
|
|
|
if qkv_bias is not None:
|
|
qkv_bias = qkv_bias.reshape(
|
|
qkv_bias.shape[0] * qkv_bias.shape[1] * qkv_bias.shape[2]
|
|
)
|
|
else:
|
|
assert len(qkv_weight.shape) == 2
|
|
assert qkv_weight.shape[0] * 3 == qkv_weight.shape[1]
|
|
if qkv_bias is not None:
|
|
assert len(qkv_bias.shape) == 1
|
|
assert qkv_bias.shape[0] == qkv_weight.shape[1]
|
|
|
|
if pre_layer_norm:
|
|
ln_out = ln_out.reshape(batch_size * seq_len, embed_dim)
|
|
qkv = fc(ln_out, qkv_weight)
|
|
if qkv_bias is not None:
|
|
qkv_bias_out = qkv + qkv_bias
|
|
else:
|
|
qkv_bias_out = qkv
|
|
ln_out = ln_out.reshape(batch_size, seq_len, embed_dim)
|
|
else:
|
|
query = query.reshape(batch_size * seq_len, embed_dim)
|
|
qkv = fc(query, qkv_weight)
|
|
if qkv_bias is not None:
|
|
qkv_bias_out = qkv + qkv_bias
|
|
else:
|
|
qkv_bias_out = qkv
|
|
query = query.reshape(batch_size, seq_len, embed_dim)
|
|
|
|
qkv_bias_out = qkv_bias_out.reshape(
|
|
batch_size, seq_len, 3, num_head, head_dim
|
|
)
|
|
# q*k^t
|
|
qkv_bias_out = qkv_bias_out.transpose(
|
|
(2, 0, 1, 3, 4)
|
|
) # 3, batch_size, seq_len, num_head, head_dim
|
|
qkv_bias_out = qkv_bias_out.transpose(
|
|
(0, 1, 3, 2, 4)
|
|
) # 3, batch_size, num_head, seq_len, head_dim
|
|
|
|
q = qkv_bias_out[0:1, ::]
|
|
q = q.reshape(batch_size, num_head, seq_len, head_dim)
|
|
k = qkv_bias_out[1:2, ::] # [1, batch_size, num_head, seq_len, head_dim]
|
|
k = k.reshape(batch_size, num_head, seq_len, head_dim)
|
|
v = qkv_bias_out[2::]
|
|
v = v.reshape(batch_size, num_head, seq_len, head_dim)
|
|
|
|
k = k.transpose([0, 1, 3, 2]) # [batch_size, num_head, head_dim, seq_len]
|
|
qkt = batch_matmul(q, k / np.sqrt(head_dim, dtype=np.float64))
|
|
|
|
if attn_mask is not None:
|
|
if attn_mask.dtype.name == 'int64':
|
|
attn_mask = (attn_mask.astype(qkt.dtype) - 1.0) * 1e9
|
|
else:
|
|
attn_mask = attn_mask.astype(qkt.dtype)
|
|
qkt += attn_mask
|
|
|
|
# softmax
|
|
softmax_out = softmax(qkt)
|
|
attn_heads = batch_matmul(softmax_out, v)
|
|
|
|
attn_heads = attn_heads.transpose(
|
|
(0, 2, 1, 3)
|
|
) # [batch_size, seq_len, num_head, head_dim]
|
|
|
|
# out_linear
|
|
out_linear_input = attn_heads.reshape(
|
|
batch_size, seq_len, num_head * head_dim
|
|
)
|
|
out_linear_out = fc(out_linear_input, out_linear_weight)
|
|
|
|
# bias add, dropout, residual add, layer_norm.
|
|
if out_linear_bias is not None:
|
|
out_linear_bias_out = out_linear_out + out_linear_bias
|
|
else:
|
|
out_linear_bias_out = out_linear_out
|
|
out_linear_bias_dropout_out = out_linear_bias_out
|
|
out_linear_bias_dropout_residual_out = query + out_linear_bias_dropout_out
|
|
if not pre_layer_norm:
|
|
out_linear_bias_dropout_residual_out = layer_norm(
|
|
out_linear_bias_dropout_residual_out,
|
|
True,
|
|
has_bias,
|
|
ln_2_scale,
|
|
ln_2_bias,
|
|
)
|
|
return out_linear_bias_dropout_residual_out
|
|
|
|
|
|
class TestFusedAttentionAPI(unittest.TestCase):
|
|
def setUp(self):
|
|
self.setXType()
|
|
self.setPreLn()
|
|
self.setAttnMask()
|
|
self.setBiasAttr()
|
|
self.setTransposeWAndB()
|
|
self.config()
|
|
self.generate_input_data()
|
|
|
|
self.rtol = 1e-5
|
|
# FIXME(limin29): Because there is a problem with the test precision
|
|
# on A100, atol is temporarily set to 1e-2, and it will be
|
|
# changed back after the precision problem is solved.
|
|
self.atol = 1e-2
|
|
# make sure local development precision
|
|
if "V100" in paddle.device.cuda.get_device_name():
|
|
self.atol = 1e-4
|
|
if self.x_type is np.float16:
|
|
self.atol = 1e-1
|
|
|
|
def setAttnMask(self):
|
|
self.has_attn_mask = True
|
|
|
|
def setBiasAttr(self):
|
|
self.bias_attr = None
|
|
|
|
def setTransposeWAndB(self):
|
|
self.transpose_qkv_wb = False
|
|
|
|
def setPreLn(self):
|
|
self.pre_layer_norm = False
|
|
|
|
def setXType(self):
|
|
self.x_type = np.float32
|
|
|
|
def config(self):
|
|
self.attn_mask_type = np.float64
|
|
self.training = True
|
|
self.need_weight = False
|
|
|
|
self.batch_size = 1
|
|
self.query_length = 2
|
|
self.head_dim = 2
|
|
self.num_heads = 2
|
|
self.embed_dim = self.head_dim * self.num_heads
|
|
|
|
self.dropout_prob = 0.0
|
|
self.attn_dropout_prob = 0.0
|
|
self.weight_attr = None
|
|
|
|
self.kdim, self.vdim = self.embed_dim, self.embed_dim
|
|
self.key_length, self.value_length = (
|
|
self.query_length,
|
|
self.query_length,
|
|
)
|
|
|
|
def generate_input_data(self):
|
|
self.query = np.random.rand(
|
|
self.batch_size, self.query_length, self.embed_dim
|
|
).astype(self.x_type)
|
|
if self.has_attn_mask:
|
|
self.attn_mask = np.ones(
|
|
(
|
|
self.batch_size,
|
|
self.num_heads,
|
|
self.query_length,
|
|
self.key_length,
|
|
),
|
|
dtype=self.attn_mask_type,
|
|
)
|
|
if self.attn_mask_type == np.int64:
|
|
self.attn_mask = np.tril(self.attn_mask)
|
|
elif self.attn_mask_type == np.float64:
|
|
self.attn_mask = (np.tril(self.attn_mask) - 1.0) * 1e9
|
|
else:
|
|
raise ValueError(
|
|
"'attn_mask_type' should be 'int64' or 'float64'."
|
|
)
|
|
else:
|
|
self.attn_mask = None
|
|
self.key, self.value = self.query, self.query
|
|
|
|
def run_imperative(self):
|
|
if self.has_attn_mask:
|
|
attn_mask_tensor = paddle.to_tensor(self.attn_mask)
|
|
else:
|
|
attn_mask_tensor = None
|
|
fused_attn = FusedMultiHeadAttention(
|
|
self.embed_dim,
|
|
self.num_heads,
|
|
self.dropout_prob,
|
|
self.attn_dropout_prob,
|
|
self.kdim,
|
|
self.vdim,
|
|
self.pre_layer_norm,
|
|
self.need_weight,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
transpose_qkv_wb=self.transpose_qkv_wb,
|
|
)
|
|
if self.bias_attr is not False:
|
|
qkv_bias = np.random.random(fused_attn.qkv_bias.shape).astype(
|
|
'float32'
|
|
)
|
|
fused_attn.qkv_bias.set_value(paddle.to_tensor(qkv_bias))
|
|
out = fused_attn(
|
|
paddle.to_tensor(self.query),
|
|
paddle.to_tensor(self.query),
|
|
paddle.to_tensor(self.query),
|
|
attn_mask_tensor,
|
|
)
|
|
|
|
fused_attn_qkv_bias = None
|
|
fused_attn_linear_bias = None
|
|
fused_attn_pre_ln_bias = None
|
|
fused_attn_ln_bias = None
|
|
if self.bias_attr is not False:
|
|
fused_attn_qkv_bias = fused_attn.qkv_bias.numpy()
|
|
fused_attn_linear_bias = fused_attn.linear_bias.numpy()
|
|
if self.pre_layer_norm:
|
|
fused_attn_pre_ln_bias = fused_attn.pre_ln_bias.numpy()
|
|
fused_attn_ln_bias = None
|
|
else:
|
|
fused_attn_pre_ln_bias = None
|
|
fused_attn_ln_bias = fused_attn.ln_bias.numpy()
|
|
|
|
ref_out = compute_reference(
|
|
self.pre_layer_norm,
|
|
self.query,
|
|
self.attn_mask,
|
|
fused_attn.pre_ln_scale.numpy() if self.pre_layer_norm else None,
|
|
fused_attn_pre_ln_bias,
|
|
fused_attn.ln_scale.numpy() if not self.pre_layer_norm else None,
|
|
fused_attn_ln_bias,
|
|
fused_attn.qkv_weight.numpy(),
|
|
fused_attn_qkv_bias,
|
|
fused_attn.linear_weight.numpy(),
|
|
fused_attn_linear_bias,
|
|
num_head=self.num_heads,
|
|
transpose_qkv_wb=self.transpose_qkv_wb,
|
|
)
|
|
np.testing.assert_allclose(
|
|
ref_out, out.numpy(), rtol=self.rtol, atol=self.atol
|
|
)
|
|
|
|
def run_static(self):
|
|
fused_attn = FusedMultiHeadAttention(
|
|
self.embed_dim,
|
|
self.num_heads,
|
|
self.dropout_prob,
|
|
self.attn_dropout_prob,
|
|
self.kdim,
|
|
self.vdim,
|
|
self.pre_layer_norm,
|
|
self.need_weight,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
self.weight_attr,
|
|
self.bias_attr,
|
|
transpose_qkv_wb=self.transpose_qkv_wb,
|
|
)
|
|
|
|
x = paddle.static.data(
|
|
name='X',
|
|
shape=[self.batch_size, self.query_length, self.embed_dim],
|
|
dtype=self.x_type,
|
|
)
|
|
if self.has_attn_mask:
|
|
attn_mask = paddle.static.data(
|
|
name='SrcMask',
|
|
shape=[
|
|
self.batch_size,
|
|
self.num_heads,
|
|
self.query_length,
|
|
self.key_length,
|
|
],
|
|
dtype=self.attn_mask_type,
|
|
)
|
|
final_out = fused_attn(x, x, x, attn_mask)
|
|
else:
|
|
final_out = fused_attn(x, x, x)
|
|
|
|
place = get_device_place()
|
|
exe = paddle.static.Executor(place)
|
|
exe.run(paddle.static.default_startup_program())
|
|
|
|
qkv_bias = None
|
|
linear_bias = None
|
|
ln_scale = None
|
|
ln_2_scale = None
|
|
ln_bias = None
|
|
ln_2_bias = None
|
|
if self.has_attn_mask:
|
|
if self.bias_attr is False:
|
|
if self.pre_layer_norm:
|
|
out, qkv_weight, out_linear_weight, ln_scale = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={"X": self.query, "SrcMask": self.attn_mask},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.linear_weight,
|
|
fused_attn.pre_ln_scale,
|
|
],
|
|
)
|
|
else:
|
|
out, qkv_weight, out_linear_weight, ln_2_scale = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={"X": self.query, "SrcMask": self.attn_mask},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.linear_weight,
|
|
fused_attn.ln_scale,
|
|
],
|
|
)
|
|
else:
|
|
if self.pre_layer_norm:
|
|
(
|
|
out,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
out_linear_weight,
|
|
linear_bias,
|
|
ln_scale,
|
|
ln_bias,
|
|
) = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={"X": self.query, "SrcMask": self.attn_mask},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.qkv_bias,
|
|
fused_attn.linear_weight,
|
|
fused_attn.linear_bias,
|
|
fused_attn.pre_ln_scale,
|
|
fused_attn.pre_ln_bias,
|
|
],
|
|
)
|
|
else:
|
|
(
|
|
out,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
out_linear_weight,
|
|
linear_bias,
|
|
ln_2_scale,
|
|
ln_2_bias,
|
|
) = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={"X": self.query, "SrcMask": self.attn_mask},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.qkv_bias,
|
|
fused_attn.linear_weight,
|
|
fused_attn.linear_bias,
|
|
fused_attn.ln_scale,
|
|
fused_attn.ln_bias,
|
|
],
|
|
)
|
|
else:
|
|
if self.bias_attr is False:
|
|
if self.pre_layer_norm:
|
|
out, qkv_weight, out_linear_weight, ln_scale = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={
|
|
"X": self.query,
|
|
},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.linear_weight,
|
|
fused_attn.pre_ln_scale,
|
|
],
|
|
)
|
|
else:
|
|
out, qkv_weight, out_linear_weight, ln_2_scale = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={
|
|
"X": self.query,
|
|
},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.linear_weight,
|
|
fused_attn.ln_scale,
|
|
],
|
|
)
|
|
else:
|
|
if self.pre_layer_norm:
|
|
(
|
|
out,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
out_linear_weight,
|
|
linear_bias,
|
|
ln_scale,
|
|
ln_bias,
|
|
) = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={
|
|
"X": self.query,
|
|
},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.qkv_bias,
|
|
fused_attn.linear_weight,
|
|
fused_attn.linear_bias,
|
|
fused_attn.pre_ln_scale,
|
|
fused_attn.pre_ln_bias,
|
|
],
|
|
)
|
|
else:
|
|
(
|
|
out,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
out_linear_weight,
|
|
linear_bias,
|
|
ln_2_scale,
|
|
ln_2_bias,
|
|
) = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={
|
|
"X": self.query,
|
|
},
|
|
fetch_list=[
|
|
final_out,
|
|
fused_attn.qkv_weight,
|
|
fused_attn.qkv_bias,
|
|
fused_attn.linear_weight,
|
|
fused_attn.linear_bias,
|
|
fused_attn.ln_scale,
|
|
fused_attn.ln_bias,
|
|
],
|
|
)
|
|
return (
|
|
out,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
out_linear_weight,
|
|
linear_bias,
|
|
ln_scale,
|
|
ln_bias,
|
|
ln_2_scale,
|
|
ln_2_bias,
|
|
)
|
|
|
|
def test_static_api(self):
|
|
paddle.enable_static()
|
|
main = paddle.static.Program()
|
|
startup = paddle.static.Program()
|
|
with paddle.static.program_guard(main, startup):
|
|
(
|
|
out,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
linear_weight,
|
|
linear_bias,
|
|
ln_scale,
|
|
ln_bias,
|
|
ln_2_scale,
|
|
ln_2_bias,
|
|
) = self.run_static()
|
|
ref_out = compute_reference(
|
|
self.pre_layer_norm,
|
|
self.query,
|
|
self.attn_mask,
|
|
ln_scale,
|
|
ln_bias,
|
|
ln_2_scale,
|
|
ln_2_bias,
|
|
qkv_weight,
|
|
qkv_bias,
|
|
linear_weight,
|
|
linear_bias,
|
|
num_head=self.num_heads,
|
|
transpose_qkv_wb=self.transpose_qkv_wb,
|
|
)
|
|
np.testing.assert_allclose(ref_out, out, rtol=self.rtol, atol=self.atol)
|
|
|
|
def test_dynamic_api(self):
|
|
paddle.disable_static(place=get_device_place())
|
|
self.run_imperative()
|
|
|
|
|
|
class TestFusedAttentionAPINoneAttnMask(TestFusedAttentionAPI):
|
|
def setAttnMask(self):
|
|
self.has_attn_mask = False
|
|
|
|
def setPreLn(self):
|
|
self.pre_layer_norm = True
|
|
|
|
|
|
class TestFusedAttentionAPIBiasIsNone(TestFusedAttentionAPI):
|
|
def setBiasAttr(self):
|
|
self.bias_attr = False
|
|
|
|
|
|
class TestFusedAttentionAPITransposeWAndB(TestFusedAttentionAPI):
|
|
def setTransposeWAndB(self):
|
|
self.transpose_qkv_wb = True
|
|
|
|
|
|
class TestFusedAttentionAPITransposeWAndBWithoutBias(TestFusedAttentionAPI):
|
|
def setTransposeWAndB(self):
|
|
self.transpose_qkv_wb = True
|
|
|
|
def setBiasAttr(self):
|
|
self.bias_attr = False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|