wehub-resource-sync
540 lines
18 KiB
Python
540 lines
18 KiB
Python
# Copyright (c) 2025 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 pprint
|
|
import random
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.compat.nn.transformer import MultiheadAttention
|
|
|
|
is_bf16_supported = (
|
|
paddle.is_compiled_with_cuda()
|
|
and paddle.cuda.get_device_capability()[0] >= 8
|
|
)
|
|
|
|
|
|
class ReferenceImplementation:
|
|
@staticmethod
|
|
def softmax(x, axis=-1):
|
|
x_max = np.max(x, axis=axis, keepdims=True)
|
|
x_max[x_max == float('-inf')] = 0.0
|
|
exp_x = np.exp(x - x_max)
|
|
sum_exp = np.sum(exp_x, axis=axis, keepdims=True)
|
|
out = exp_x / (sum_exp + 1e-10)
|
|
return out
|
|
|
|
@staticmethod
|
|
def linear(x, weight, bias=None):
|
|
res = x @ weight
|
|
if bias is not None:
|
|
res += bias
|
|
return res
|
|
|
|
@staticmethod
|
|
def forward(
|
|
query,
|
|
key,
|
|
value,
|
|
w_q,
|
|
w_k,
|
|
w_v,
|
|
w_out,
|
|
b_q,
|
|
b_k,
|
|
b_v,
|
|
b_out,
|
|
bias_k=None,
|
|
bias_v=None,
|
|
key_padding_mask=None,
|
|
attn_mask=None,
|
|
add_bias_kv=False,
|
|
add_zero_attn=False,
|
|
num_heads=4,
|
|
need_weights=True,
|
|
average_attn_weights=True,
|
|
):
|
|
is_batched = query.ndim == 3
|
|
if not is_batched:
|
|
query = query.reshape([1, *query.shape])
|
|
key = key.reshape([1, *key.shape])
|
|
value = value.reshape([1, *value.shape])
|
|
B, L, E = query.shape
|
|
|
|
head_dim = E // num_heads
|
|
scale = head_dim**-0.5
|
|
|
|
q = ReferenceImplementation.linear(query, w_q, b_q)
|
|
k = ReferenceImplementation.linear(key, w_k, b_k)
|
|
v = ReferenceImplementation.linear(value, w_v, b_v)
|
|
|
|
pad_col_count = 0
|
|
if add_bias_kv:
|
|
if bias_k is not None:
|
|
bk = np.tile(bias_k, (B, 1, 1))
|
|
bv = np.tile(bias_v, (B, 1, 1))
|
|
k = np.concatenate([k, bk], axis=1)
|
|
v = np.concatenate([v, bv], axis=1)
|
|
pad_col_count += 1
|
|
|
|
if add_zero_attn:
|
|
zeros = np.zeros((B, 1, E), dtype=q.dtype)
|
|
k = np.concatenate([k, zeros], axis=1)
|
|
v = np.concatenate([v, zeros], axis=1)
|
|
pad_col_count += 1
|
|
|
|
curr_S = k.shape[1]
|
|
|
|
q = q.reshape(B, L, num_heads, head_dim).transpose(0, 2, 1, 3)
|
|
k = k.reshape(B, curr_S, num_heads, head_dim).transpose(0, 2, 1, 3)
|
|
v = v.reshape(B, curr_S, num_heads, head_dim).transpose(0, 2, 1, 3)
|
|
|
|
scores = np.matmul(q, k.transpose(0, 1, 3, 2))
|
|
scores = scores * scale
|
|
|
|
def pad_mask_width(mask_arr, pad_amt):
|
|
if pad_amt == 0:
|
|
return mask_arr
|
|
shape = list(mask_arr.shape)
|
|
shape[-1] = pad_amt
|
|
if mask_arr.dtype == bool:
|
|
pad = np.zeros(shape, dtype=bool)
|
|
else:
|
|
pad = np.zeros(shape, dtype=mask_arr.dtype)
|
|
return np.concatenate([mask_arr, pad], axis=-1)
|
|
|
|
if attn_mask is not None:
|
|
am = attn_mask
|
|
if pad_col_count > 0:
|
|
am = pad_mask_width(am, pad_col_count)
|
|
|
|
if am.ndim == 2:
|
|
am = am[None, None, :, :]
|
|
elif am.ndim == 3:
|
|
if am.shape[0] == B * num_heads:
|
|
am = am.reshape(B, num_heads, L, -1)
|
|
elif am.shape[0] == B:
|
|
am = am[:, None, :, :]
|
|
|
|
if am.dtype == bool:
|
|
scores = np.where(am, float('-inf'), scores)
|
|
else:
|
|
scores += am
|
|
|
|
if key_padding_mask is not None:
|
|
kpm = key_padding_mask
|
|
if pad_col_count > 0:
|
|
kpm = pad_mask_width(kpm, pad_col_count)
|
|
kpm = kpm[:, None, None, :]
|
|
|
|
if kpm.dtype == bool:
|
|
scores = np.where(kpm, float('-inf'), scores)
|
|
else:
|
|
scores += kpm
|
|
|
|
attn_weights = ReferenceImplementation.softmax(scores, axis=-1)
|
|
|
|
ctx = np.matmul(attn_weights, v)
|
|
ctx = ctx.transpose(0, 2, 1, 3).reshape(B, L, E)
|
|
output = ReferenceImplementation.linear(ctx, w_out, b_out)
|
|
|
|
if need_weights:
|
|
if average_attn_weights:
|
|
attn_weights = np.mean(attn_weights, axis=1)
|
|
else:
|
|
attn_weights = None
|
|
if not is_batched:
|
|
output = output.reshape(output.shape[1:])
|
|
if attn_weights is not None:
|
|
attn_weights = attn_weights.reshape(attn_weights.shape[1:])
|
|
return output, attn_weights
|
|
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_cuda(),
|
|
"SDPA is not fully supported on non-CUDA devices.",
|
|
)
|
|
class TestMHA_Coverage(unittest.TestCase):
|
|
def setUp(self):
|
|
self.seed = 42
|
|
self.random_seed()
|
|
self.atol = 1e-3
|
|
self.num_fuzz_iter = 200
|
|
|
|
def random_seed(self):
|
|
random.seed(self.seed)
|
|
np.random.seed(self.seed)
|
|
paddle.seed(self.seed)
|
|
|
|
def _extract_weights(self, layer):
|
|
sd = layer.state_dict()
|
|
w = {}
|
|
|
|
def to_np(t):
|
|
return t.cast('float32').numpy() if t is not None else None
|
|
|
|
def safe_T(arr):
|
|
return arr.T if arr is not None else None
|
|
|
|
w['w_out'] = safe_T(to_np(sd.get('out_proj.weight')))
|
|
w['b_out'] = to_np(sd.get('out_proj.bias'))
|
|
w['bias_k'] = to_np(sd.get('bias_k'))
|
|
w['bias_v'] = to_np(sd.get('bias_v'))
|
|
|
|
if layer._qkv_same_embed_dim:
|
|
in_w = to_np(sd.get('in_proj_weight'))
|
|
if in_w is not None:
|
|
in_w_t = in_w.T
|
|
w['w_q'], w['w_k'], w['w_v'] = np.split(in_w_t, 3, axis=1)
|
|
else:
|
|
w['w_q'] = w['w_k'] = w['w_v'] = None
|
|
|
|
if sd.get('in_proj_bias') is not None:
|
|
in_b = to_np(sd['in_proj_bias'])
|
|
w['b_q'], w['b_k'], w['b_v'] = np.split(in_b, 3, axis=0)
|
|
else:
|
|
w['b_q'] = w['b_k'] = w['b_v'] = None
|
|
else:
|
|
w['w_q'] = safe_T(to_np(sd.get('q_proj_weight')))
|
|
w['w_k'] = safe_T(to_np(sd.get('k_proj_weight')))
|
|
w['w_v'] = safe_T(to_np(sd.get('v_proj_weight')))
|
|
w['b_q'] = to_np(sd.get('q_proj_bias'))
|
|
w['b_k'] = to_np(sd.get('k_proj_bias'))
|
|
w['b_v'] = to_np(sd.get('v_proj_bias'))
|
|
|
|
return w
|
|
|
|
def generate_config(self, **overrides):
|
|
"""Generates a complete configuration dict, using defaults or random values where needed."""
|
|
config = {}
|
|
|
|
# Basic Dimensions
|
|
config['num_heads'] = overrides.get(
|
|
'num_heads', random.choice([1, 2, 4])
|
|
)
|
|
default_embed_dim = random.randint(4, 12) * config['num_heads']
|
|
config['embed_dim'] = overrides.get('embed_dim', default_embed_dim)
|
|
|
|
config['B'] = overrides.get('B', random.randint(1, 4))
|
|
config['L'] = overrides.get('L', random.randint(2, 8))
|
|
|
|
# Cross Attention Logic
|
|
config['is_cross'] = overrides.get(
|
|
'is_cross', random.choice([True, False])
|
|
)
|
|
|
|
if not config['is_cross']:
|
|
config['S'] = config['L']
|
|
else:
|
|
config['S'] = overrides.get('S', random.randint(2, 8))
|
|
|
|
# Key/Value Dimensions
|
|
if not config['is_cross']:
|
|
config['kdim'] = config['embed_dim']
|
|
config['vdim'] = config['embed_dim']
|
|
else:
|
|
config['kdim'] = overrides.get('kdim', random.randint(4, 12))
|
|
config['vdim'] = overrides.get('vdim', random.randint(4, 12))
|
|
|
|
# Booleans
|
|
config['batch_first'] = overrides.get(
|
|
'batch_first', random.choice([True, False])
|
|
)
|
|
config['bias'] = overrides.get('bias', random.choice([True, False]))
|
|
config['dtype'] = overrides.get('dtype', 'float32')
|
|
config['need_weights'] = overrides.get('need_weights', True)
|
|
config['add_bias_kv'] = overrides.get('add_bias_kv', False)
|
|
config['add_zero_attn'] = overrides.get('add_zero_attn', False)
|
|
config['average_attn_weights'] = overrides.get(
|
|
'average_attn_weights', random.choice([True, False])
|
|
)
|
|
config['is_causal'] = overrides.get('is_causal', False)
|
|
|
|
# Special flags
|
|
config['key_padding_mask'] = overrides.get(
|
|
'key_padding_mask', random.random() < 0.5
|
|
)
|
|
|
|
# Unbatched input simulation (B=1 case)
|
|
# If B=1, we sometimes pass 2D inputs [L, D] instead of [1, L, D]
|
|
if config['B'] == 1:
|
|
config['unbatched_input'] = overrides.get(
|
|
'unbatched_input', random.random() < 0.5
|
|
)
|
|
else:
|
|
config['unbatched_input'] = False
|
|
|
|
config['random_mask'] = overrides.get(
|
|
'random_mask', random.random() < 0.5
|
|
)
|
|
config['random_mask_3d'] = overrides.get(
|
|
'random_mask_3d', random.random() < 0.5
|
|
)
|
|
|
|
if config['random_mask']:
|
|
config['is_causal'] = False
|
|
|
|
return config
|
|
|
|
def run_case(self, config):
|
|
B = config['B']
|
|
L = config['L']
|
|
S = config['S']
|
|
H = config['num_heads']
|
|
D = config['embed_dim']
|
|
|
|
pd_dtype = getattr(paddle, config['dtype'])
|
|
|
|
model = MultiheadAttention(
|
|
embed_dim=D,
|
|
num_heads=H,
|
|
dropout=0.0,
|
|
bias=config['bias'],
|
|
batch_first=config['batch_first'],
|
|
kdim=config['kdim'],
|
|
vdim=config['vdim'],
|
|
add_bias_kv=config['add_bias_kv'],
|
|
add_zero_attn=config['add_zero_attn'],
|
|
dtype=pd_dtype,
|
|
)
|
|
model.eval()
|
|
|
|
q_shape = [B, L, D] if config['batch_first'] else [L, B, D]
|
|
k_shape = (
|
|
[B, S, config['kdim']]
|
|
if config['batch_first']
|
|
else [S, B, config['kdim']]
|
|
)
|
|
v_shape = (
|
|
[B, S, config['vdim']]
|
|
if config['batch_first']
|
|
else [S, B, config['vdim']]
|
|
)
|
|
|
|
if config['unbatched_input']:
|
|
q_shape = [L, D]
|
|
k_shape = [S, config['kdim']]
|
|
v_shape = [S, config['vdim']]
|
|
|
|
q_pd = paddle.randn(q_shape).cast(pd_dtype)
|
|
k_pd = (
|
|
paddle.randn(k_shape).cast(pd_dtype) if config['is_cross'] else q_pd
|
|
)
|
|
v_pd = (
|
|
paddle.randn(v_shape).cast(pd_dtype) if config['is_cross'] else q_pd
|
|
)
|
|
|
|
attn_mask = None
|
|
key_padding_mask = None
|
|
|
|
if config['is_causal'] and config['random_mask']:
|
|
raise ValueError(
|
|
"Both is_causal and random_mask cannot be True at the same time."
|
|
)
|
|
|
|
if config['is_causal']:
|
|
attn_mask = np.triu(np.ones((L, S), dtype="bool"), k=1)
|
|
attn_mask = paddle.to_tensor(attn_mask)
|
|
elif config['random_mask']:
|
|
if config['random_mask_3d']:
|
|
mask_vals = np.random.choice(
|
|
[True, False], size=(B * H, L, S), p=[0.2, 0.8]
|
|
)
|
|
mask_vals[:, :, 0] = False
|
|
attn_mask = paddle.to_tensor(mask_vals)
|
|
else:
|
|
mask_vals = np.random.choice(
|
|
[True, False], size=(L, S), p=[0.2, 0.8]
|
|
)
|
|
mask_vals[
|
|
:,
|
|
0,
|
|
] = False
|
|
attn_mask = paddle.to_tensor(mask_vals)
|
|
|
|
if config['key_padding_mask']:
|
|
kp_np = np.random.choice([True, False], size=(B, S), p=[0.2, 0.8])
|
|
kp_np[:, 0] = False
|
|
key_padding_mask = paddle.to_tensor(kp_np)
|
|
|
|
with paddle.no_grad():
|
|
out_pd, w_pd = model(
|
|
q_pd,
|
|
k_pd,
|
|
v_pd,
|
|
key_padding_mask=key_padding_mask,
|
|
attn_mask=attn_mask,
|
|
need_weights=config['need_weights'],
|
|
average_attn_weights=config['average_attn_weights'],
|
|
is_causal=config['is_causal'],
|
|
)
|
|
|
|
q_np = q_pd.cast('float32').numpy()
|
|
k_np = k_pd.cast('float32').numpy()
|
|
v_np = v_pd.cast('float32').numpy()
|
|
|
|
if not config['batch_first'] and len(q_np.shape) == 3:
|
|
q_np = q_np.transpose(1, 0, 2)
|
|
k_np = k_np.transpose(1, 0, 2)
|
|
v_np = v_np.transpose(1, 0, 2)
|
|
|
|
weights = self._extract_weights(model)
|
|
|
|
kp_np = (
|
|
key_padding_mask.numpy() if key_padding_mask is not None else None
|
|
)
|
|
am_np = attn_mask.numpy() if attn_mask is not None else None
|
|
|
|
out_ref, w_ref = ReferenceImplementation.forward(
|
|
q_np,
|
|
k_np,
|
|
v_np,
|
|
w_q=weights['w_q'],
|
|
w_k=weights['w_k'],
|
|
w_v=weights['w_v'],
|
|
w_out=weights['w_out'],
|
|
b_q=weights['b_q'],
|
|
b_k=weights['b_k'],
|
|
b_v=weights['b_v'],
|
|
b_out=weights['b_out'],
|
|
bias_k=weights['bias_k'],
|
|
bias_v=weights['bias_v'],
|
|
key_padding_mask=kp_np,
|
|
attn_mask=am_np,
|
|
add_bias_kv=config['add_bias_kv'],
|
|
add_zero_attn=config['add_zero_attn'],
|
|
num_heads=H,
|
|
need_weights=config['need_weights'],
|
|
average_attn_weights=config['average_attn_weights'],
|
|
)
|
|
|
|
if len(q_np.shape) == 3 and not config['batch_first']:
|
|
out_ref = out_ref.transpose(1, 0, 2)
|
|
|
|
current_atol = 1e-3 if config['dtype'] == 'float16' else self.atol
|
|
if not config['need_weights'] and config['dtype'] == 'float16':
|
|
current_atol = 1e-3
|
|
|
|
if not paddle.is_compiled_with_custom_device("dcu"):
|
|
current_atol = 1e-2
|
|
|
|
# Pretty print config for error message
|
|
config_str = pprint.pformat(config)
|
|
try:
|
|
np.testing.assert_allclose(
|
|
out_pd.cast('float32').numpy(),
|
|
out_ref,
|
|
atol=current_atol,
|
|
rtol=current_atol,
|
|
err_msg=f"\nOutput mismatch.\nConfig:\n{config_str}",
|
|
)
|
|
except AssertionError as e:
|
|
print(f"Failed with config: {config}")
|
|
out_ref, w_ref = ReferenceImplementation.forward(
|
|
q_np,
|
|
k_np,
|
|
v_np,
|
|
w_q=weights['w_q'],
|
|
w_k=weights['w_k'],
|
|
w_v=weights['w_v'],
|
|
w_out=weights['w_out'],
|
|
b_q=weights['b_q'],
|
|
b_k=weights['b_k'],
|
|
b_v=weights['b_v'],
|
|
b_out=weights['b_out'],
|
|
bias_k=weights['bias_k'],
|
|
bias_v=weights['bias_v'],
|
|
key_padding_mask=kp_np,
|
|
attn_mask=am_np,
|
|
add_bias_kv=config['add_bias_kv'],
|
|
add_zero_attn=config['add_zero_attn'],
|
|
num_heads=H,
|
|
need_weights=config['need_weights'],
|
|
average_attn_weights=config['average_attn_weights'],
|
|
)
|
|
out_pd, w_pd = model(
|
|
q_pd,
|
|
k_pd,
|
|
v_pd,
|
|
key_padding_mask=key_padding_mask,
|
|
attn_mask=attn_mask,
|
|
need_weights=config['need_weights'],
|
|
average_attn_weights=config['average_attn_weights'],
|
|
is_causal=config['is_causal'],
|
|
)
|
|
raise e
|
|
|
|
if config['need_weights'] and w_pd is not None:
|
|
np.testing.assert_allclose(
|
|
w_pd.cast('float32').numpy(),
|
|
w_ref,
|
|
atol=current_atol,
|
|
rtol=current_atol,
|
|
err_msg=f"\nWeights mismatch.\nConfig:\n{config_str}",
|
|
)
|
|
|
|
if not config['need_weights']:
|
|
self.assertIsNone(w_pd)
|
|
|
|
def test_add_bias_kv(self):
|
|
config = self.generate_config(add_bias_kv=True, add_zero_attn=False)
|
|
self.run_case(config)
|
|
|
|
def test_add_zero_attn(self):
|
|
config = self.generate_config(add_bias_kv=False, add_zero_attn=True)
|
|
self.run_case(config)
|
|
|
|
def test_bias_kv_and_zero_attn(self):
|
|
config = self.generate_config(add_bias_kv=True, add_zero_attn=True)
|
|
self.run_case(config)
|
|
|
|
def test_is_causal(self):
|
|
config = self.generate_config(is_causal=True, is_cross=False)
|
|
self.run_case(config)
|
|
|
|
def test_sdpa_path(self):
|
|
if not paddle.is_compiled_with_cuda():
|
|
return
|
|
|
|
config = self.generate_config(
|
|
dtype='float16', need_weights=False, num_heads=16, embed_dim=32
|
|
)
|
|
try:
|
|
self.run_case(config)
|
|
except AssertionError:
|
|
pass
|
|
|
|
def test_random_fuzz(self):
|
|
for _ in range(self.num_fuzz_iter):
|
|
config = self.generate_config(
|
|
add_bias_kv=random.choice([True, False]),
|
|
add_zero_attn=random.choice([True, False]),
|
|
need_weights=random.choice([True, False]),
|
|
is_causal=random.choice([True, False]),
|
|
dtype=(
|
|
random.choice(['float32', 'bfloat16', 'float16'])
|
|
if is_bf16_supported
|
|
else 'float32'
|
|
),
|
|
random_mask=random.choice([True, False]),
|
|
random_mask_3d=random.choice([True, False]),
|
|
)
|
|
self.run_case(config)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|