wehub-resource-sync
494 lines
15 KiB
Python
494 lines
15 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.
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import random
|
|
import unittest
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
from op_test import get_cuda_version, get_device_place, is_custom_device
|
|
|
|
import paddle
|
|
import paddle.incubate.nn.attn_bias as ab
|
|
import paddle.nn.functional as F
|
|
from paddle.base import core
|
|
from paddle.incubate.nn.memory_efficient_attention import (
|
|
memory_efficient_attention,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
paddle.seed(2023)
|
|
|
|
|
|
def create_attn_bias(
|
|
bias_type,
|
|
batch_size: int,
|
|
num_heads: int,
|
|
q_len: int,
|
|
kv_len: int,
|
|
tdtype,
|
|
pdtype,
|
|
requires_grad: bool,
|
|
fmt: str,
|
|
):
|
|
if bias_type is None or isinstance(None, bias_type):
|
|
return None
|
|
r = random.Random(
|
|
"-".join(map(str, [batch_size, q_len, kv_len, tdtype, fmt]))
|
|
)
|
|
if bias_type is paddle.Tensor:
|
|
if fmt == "BMK":
|
|
batch_size *= num_heads
|
|
num_heads = 1
|
|
attn_bias = (
|
|
paddle.randn((batch_size, num_heads, 1, kv_len), dtype=pdtype) * 3
|
|
)
|
|
attn_bias = attn_bias.expand([batch_size, num_heads, q_len, kv_len])
|
|
if requires_grad:
|
|
attn_bias.stop_gradient = False
|
|
return attn_bias
|
|
if bias_type is ab.LowerTriangularMask:
|
|
return ab.LowerTriangularMask()
|
|
if bias_type in [
|
|
ab.BlockDiagonalMask,
|
|
ab.BlockDiagonalCausalMask,
|
|
]:
|
|
# This bias is not supported in BMK format
|
|
assert fmt == "BMHK"
|
|
block_diag = ab.BlockDiagonalMask.from_seqlens(
|
|
*_rand_seqlens(r, batch_size, q_len, kv_len)
|
|
)
|
|
if bias_type is ab.BlockDiagonalCausalMask:
|
|
block_diag = block_diag.make_causal()
|
|
return block_diag
|
|
raise AssertionError(f"Unsupported bias type: {bias_type}")
|
|
|
|
|
|
def _rand_seqlens(
|
|
r: random.Random, bs: int, q_len: int, kv_len: int
|
|
) -> tuple[Sequence[int], Sequence[int]]:
|
|
q_len *= bs
|
|
kv_len *= bs
|
|
seqlens_q: list[int] = []
|
|
seqlens_k: list[int] = []
|
|
|
|
step_q = [max(1, q_len // 10), max(2, q_len // 2)]
|
|
step_k = [max(1, kv_len // 10), max(2, kv_len // 2)]
|
|
while sum(seqlens_q) < q_len and sum(seqlens_k) < kv_len:
|
|
seqlens_q.append(r.randrange(*step_q))
|
|
seqlens_k.append(r.randrange(*step_k))
|
|
seqlens_q[-1] = q_len - sum(seqlens_q[:-1])
|
|
seqlens_k[-1] = kv_len - sum(seqlens_k[:-1])
|
|
return seqlens_q, seqlens_k
|
|
|
|
|
|
def attention_naive(q, k, v, attn_bias, dropout_prob, scale, seed):
|
|
qt = paddle.transpose(q, [0, 2, 1, 3])
|
|
kt = paddle.transpose(k, [0, 2, 1, 3])
|
|
vt = paddle.transpose(v, [0, 2, 1, 3])
|
|
scale = 1.0 / np.sqrt(q.shape[-1])
|
|
s = paddle.matmul(qt, paddle.transpose(kt, [0, 1, 3, 2]))
|
|
s = paddle.scale(s, scale)
|
|
|
|
if attn_bias is None:
|
|
dropout_input = F.softmax(s)
|
|
elif isinstance(
|
|
attn_bias,
|
|
(
|
|
ab.LowerTriangularMask,
|
|
ab.BlockDiagonalMask,
|
|
ab.BlockDiagonalCausalMask,
|
|
),
|
|
):
|
|
bias = attn_bias.materialize(
|
|
(q.shape[0], q.shape[2], q.shape[1], k.shape[1]), q.dtype
|
|
)
|
|
dropout_input = F.softmax(s + bias)
|
|
elif isinstance(attn_bias, paddle.Tensor):
|
|
dropout_input = F.softmax(s + attn_bias)
|
|
|
|
paddle.seed(seed)
|
|
dropout_output = F.dropout(
|
|
x=dropout_input,
|
|
p=dropout_prob,
|
|
training=True,
|
|
mode="upscale_in_train",
|
|
)
|
|
|
|
o = paddle.matmul(dropout_output, vt)
|
|
return paddle.transpose(o, [0, 2, 1, 3])
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11030,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.3",
|
|
)
|
|
class TestMemEffAttentionAPI(unittest.TestCase):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 128, 8, 16)
|
|
self.dtype = 'float32'
|
|
self.dropout = 0.0
|
|
self.training = True
|
|
self.attention_bias = None
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
def test_all(self):
|
|
print(
|
|
f"Test All case shape {self.shape} dtype {self.dtype} name {self.name}"
|
|
)
|
|
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
q_ = paddle.to_tensor(
|
|
query, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
k_ = paddle.to_tensor(
|
|
key, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
v_ = paddle.to_tensor(
|
|
value, place=self.place, dtype=self.dtype, stop_gradient=False
|
|
)
|
|
|
|
q.stop_gradient = False
|
|
k.stop_gradient = False
|
|
v.stop_gradient = False
|
|
q_.stop_gradient = False
|
|
k_.stop_gradient = False
|
|
v_.stop_gradient = False
|
|
|
|
out_ = attention_naive(
|
|
q_, k_, v_, self.attention_bias, self.dropout, self.scale, self.seed
|
|
)
|
|
|
|
paddle.seed(self.seed)
|
|
out = memory_efficient_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
self.attention_bias,
|
|
self.dropout,
|
|
self.scale,
|
|
self.training,
|
|
)
|
|
|
|
np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03)
|
|
|
|
grad_out = paddle.ones_like(q)
|
|
|
|
out.backward(grad_out)
|
|
out_.backward(grad_out)
|
|
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-03, atol=1e-03
|
|
)
|
|
|
|
|
|
class TestMemEffAPIDtypeFp16(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp16"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 128, 128)
|
|
self.dtype = paddle.float16
|
|
self.dropout = 0.0
|
|
self.attention_bias = None
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIShape0(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 128, 32)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = None
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIShape1(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 16, 16)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = None
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIShape2(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 8, 8)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = None
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIShape3(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32"
|
|
self.place = get_device_place()
|
|
self.shape = (16, 32, 128, 128)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = None
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIMask0(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32_BlockDiagonalMask"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 128, 128)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = create_attn_bias(
|
|
ab.BlockDiagonalMask,
|
|
self.shape[0],
|
|
self.shape[2],
|
|
self.shape[1],
|
|
self.shape[1],
|
|
"float32",
|
|
self.dtype,
|
|
False,
|
|
"BMHK",
|
|
)
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIMask1(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32_BlockDiagonalCausalMask"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 128, 128)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = create_attn_bias(
|
|
ab.BlockDiagonalCausalMask,
|
|
self.shape[0],
|
|
self.shape[2],
|
|
self.shape[1],
|
|
self.shape[1],
|
|
"float32",
|
|
self.dtype,
|
|
False,
|
|
"BMHK",
|
|
)
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIMask2(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32_LowerTriangularMask"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 128, 128)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = create_attn_bias(
|
|
ab.LowerTriangularMask,
|
|
self.shape[0],
|
|
self.shape[2],
|
|
self.shape[1],
|
|
self.shape[1],
|
|
"float32",
|
|
self.dtype,
|
|
False,
|
|
"BMHK",
|
|
)
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
class TestMemEffAPIMask3(TestMemEffAttentionAPI):
|
|
def setUp(self):
|
|
self.name = "MemEffAPI_fp32_AnyTensor"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 32, 128, 128)
|
|
self.dtype = paddle.float32
|
|
self.dropout = 0.0
|
|
self.attention_bias = (
|
|
paddle.randn(
|
|
(self.shape[0], self.shape[2], 1, self.shape[1]),
|
|
dtype=self.dtype,
|
|
)
|
|
* 3
|
|
)
|
|
self.attention_bias = self.attention_bias.expand(
|
|
[self.shape[0], self.shape[2], self.shape[1], self.shape[1]]
|
|
)
|
|
self.attention_bias.stop_gradient = False
|
|
self.training = True
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device())
|
|
or get_cuda_version() < 11030,
|
|
"core is not compiled with CUDA and cuda version need larger than or equal to 11.3",
|
|
)
|
|
class TestMemEffAttentionAPIWithStopGradient(unittest.TestCase):
|
|
def setUp(self):
|
|
self.name = "MemEffAttnQKV_FFF"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 128, 8, 16)
|
|
self.dtype = 'float32'
|
|
self.dropout = 0.0
|
|
self.training = True
|
|
self.attention_bias = None
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
self.q_grad_stop_gradient = True
|
|
self.k_grad_stop_gradient = False
|
|
self.v_grad_stop_gradient = False
|
|
|
|
def test_all(self):
|
|
logging.info(
|
|
f"Test All case shape {self.shape} dtype {self.dtype} name {self.name}"
|
|
)
|
|
|
|
paddle.disable_static()
|
|
|
|
query = np.random.random(self.shape)
|
|
q = paddle.to_tensor(
|
|
query,
|
|
place=self.place,
|
|
dtype=self.dtype,
|
|
stop_gradient=self.q_grad_stop_gradient,
|
|
)
|
|
q_ = paddle.to_tensor(
|
|
query,
|
|
place=self.place,
|
|
dtype=self.dtype,
|
|
stop_gradient=self.q_grad_stop_gradient,
|
|
)
|
|
key = np.random.random(self.shape)
|
|
k = paddle.to_tensor(
|
|
key,
|
|
place=self.place,
|
|
dtype=self.dtype,
|
|
stop_gradient=self.k_grad_stop_gradient,
|
|
)
|
|
k_ = paddle.to_tensor(
|
|
key,
|
|
place=self.place,
|
|
dtype=self.dtype,
|
|
stop_gradient=self.k_grad_stop_gradient,
|
|
)
|
|
value = np.random.random(self.shape)
|
|
v = paddle.to_tensor(
|
|
value,
|
|
place=self.place,
|
|
dtype=self.dtype,
|
|
stop_gradient=self.v_grad_stop_gradient,
|
|
)
|
|
v_ = paddle.to_tensor(
|
|
value,
|
|
place=self.place,
|
|
dtype=self.dtype,
|
|
stop_gradient=self.v_grad_stop_gradient,
|
|
)
|
|
|
|
out_ = attention_naive(
|
|
q_, k_, v_, self.attention_bias, self.dropout, self.scale, self.seed
|
|
)
|
|
|
|
paddle.seed(self.seed)
|
|
out = memory_efficient_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
self.attention_bias,
|
|
self.dropout,
|
|
self.scale,
|
|
self.training,
|
|
)
|
|
|
|
np.testing.assert_allclose(out.numpy(), out_, rtol=5e-03, atol=1e-03)
|
|
|
|
out.backward()
|
|
out_.backward()
|
|
|
|
if q.stop_gradient is not True:
|
|
np.testing.assert_allclose(
|
|
q.grad.numpy(), q_.grad.numpy(), rtol=5e-03, atol=1e-03
|
|
)
|
|
|
|
if k.stop_gradient is not True:
|
|
np.testing.assert_allclose(
|
|
k.grad.numpy(), k.grad.numpy(), rtol=5e-03, atol=1e-03
|
|
)
|
|
if v.stop_gradient is not True:
|
|
np.testing.assert_allclose(
|
|
v.grad.numpy(), v_.grad.numpy(), rtol=5e-03, atol=1e-03
|
|
)
|
|
|
|
|
|
class TestQKVFTT(TestMemEffAttentionAPIWithStopGradient):
|
|
def setUp(self):
|
|
self.name = "MemEffAttnQKV_TTT"
|
|
self.place = get_device_place()
|
|
self.shape = (1, 128, 8, 16)
|
|
self.dtype = 'float32'
|
|
self.dropout = 0.0
|
|
self.training = True
|
|
self.attention_bias = None
|
|
self.scale = 1.0 / np.sqrt(self.shape[-1])
|
|
self.seed = 2023
|
|
self.q_grad_stop_gradient = False
|
|
self.k_grad_stop_gradient = True
|
|
self.v_grad_stop_gradient = True
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|