已合并
Add test_native_mha.py. #21352
liu-jiaweikf创建于 2025年5月29日
Add test_native_mha.py. #21352
已合并
从refs/pull/21352/head合入到v2.7.1
共 1 个文件变更+353-0
| @@ -0,0 +1,353 @@ | |||
| 1 | +# Owner(s): ["module: nn"] | ||
| 2 | +import math | ||
| 3 | +import copy | ||
| 4 | +import unittest | ||
| 5 | + | ||
| 6 | +import torch | ||
| 7 | +import torch_npu | ||
| 8 | +import torch_npu.testing | ||
| 9 | +from torch.testing._internal.common_device_type import ( | ||
| 10 | + dtypes, | ||
| 11 | + dtypesIfPRIVATEUSE1, | ||
| 12 | + instantiate_device_type_tests, | ||
| 13 | + onlyPRIVATEUSE1, | ||
| 14 | + skipMeta, | ||
| 15 | +) | ||
| 16 | +from torch.testing._internal.common_utils import parametrize, run_tests, TestCase | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +class TestMHADeviceType(TestCase): | ||
| 20 | + | ||
| 21 | + def _test_transform_bias_rescale_qkv_impl( | ||
| 22 | + self, device, dtype, use_nt, use_padding=False | ||
| 23 | + ): | ||
| 24 | + tests = [ | ||
| 25 | + (64, 4, 16, 8), | ||
| 26 | + # dim_per_head = 12 does not divide evenly by CPU vectorization length of 8 | ||
| 27 | + (24, 2, 4, 2), | ||
| 28 | + # Make sure CUDA can handle small input sizes | ||
| 29 | + (2, 2, 2, 2), | ||
| 30 | + # dim_per_head = 6 does not divide evenly by CUDA vectorization length of 4, | ||
| 31 | + # causes alignment issues | ||
| 32 | + (24, 4, 4, 2), | ||
| 33 | + (48, 4, 16, 8), | ||
| 34 | + ] | ||
| 35 | + for (embed_dim, num_heads, bs, sl) in tests: | ||
| 36 | + with self.subTest(embed_dim=embed_dim, num_heads=num_heads, bs=bs, sl=sl): | ||
| 37 | + torch.manual_seed(9343) | ||
| 38 | + dense_x = x = ( | ||
| 39 | + torch.randn(bs, sl, 3 * embed_dim, device=device, dtype=dtype) * 10 | ||
| 40 | + ) | ||
| 41 | + if use_padding: | ||
| 42 | + x[0][-1] = torch.full(x[0][-1].shape, float("-Inf")) | ||
| 43 | + if use_nt: | ||
| 44 | + xs = list(torch.unbind(x)) | ||
| 45 | + if use_padding: | ||
| 46 | + xs[0] = xs[0][:-1] | ||
| 47 | + x = torch.nested.nested_tensor(xs, device=device, dtype=dtype) | ||
| 48 | + qkv = torch.nn.Linear(embed_dim, 3 * embed_dim, device=device, dtype=dtype) | ||
| 49 | + | ||
| 50 | + # We have to use inference_mode here because q/k/v are | ||
| 51 | + # all views of the same Tensor, which autograd doesn't | ||
| 52 | + # like. This is fine because this function is only | ||
| 53 | + # exposed to Python for purposes of writing this test. | ||
| 54 | + with torch.inference_mode(): | ||
| 55 | + (q, k, v) = torch._transform_bias_rescale_qkv( | ||
| 56 | + x, qkv.bias, num_heads=num_heads | ||
| 57 | + ) | ||
| 58 | + | ||
| 59 | + def simple_transform_bias_rescale_qkv(qkv, bias): | ||
| 60 | + (q, k, v) = torch.split(qkv, embed_dim, dim=-1) | ||
| 61 | + (q_bias, k_bias, v_bias) = torch.split(bias, embed_dim, dim=-1) | ||
| 62 | + | ||
| 63 | + def embiggen(x): | ||
| 64 | + if not use_nt: | ||
| 65 | + return x | ||
| 66 | + b, t, d = x.size() | ||
| 67 | + t = t + (8 - t % 8) % 8 | ||
| 68 | + newsize = (b, t, d) | ||
| 69 | + new_x = torch.zeros(newsize, device=device, dtype=dtype) | ||
| 70 | + new_x[:x.size()[0], :x.size()[1], :x.size()[2]] = x | ||
| 71 | + return new_x | ||
| 72 | + return tuple( | ||
| 73 | + embiggen(x).reshape( | ||
| 74 | + (bs, -1, num_heads, embed_dim // num_heads) | ||
| 75 | + ).transpose(2, 1) | ||
| 76 | + for x in ( | ||
| 77 | + (q + q_bias) / math.sqrt(embed_dim // num_heads), | ||
| 78 | + (k + k_bias), | ||
| 79 | + (v + v_bias), | ||
| 80 | + ) | ||
| 81 | + ) | ||
| 82 | + | ||
| 83 | + correct_q, correct_k, correct_v = simple_transform_bias_rescale_qkv( | ||
| 84 | + dense_x, qkv.bias | ||
| 85 | + ) | ||
| 86 | + if use_nt and use_padding: | ||
| 87 | + for t in (correct_q, correct_k, correct_v): | ||
| 88 | + t[t == float("-Inf")] = 0 | ||
| 89 | + | ||
| 90 | + self.assertEqual(q.size(), correct_q.size()) | ||
| 91 | + torch.testing.assert_close(q, correct_q) | ||
| 92 | + torch.testing.assert_close(k, correct_k) | ||
| 93 | + torch.testing.assert_close(v, correct_v) | ||
| 94 | + | ||
| 95 | + | ||
| 96 | + | ||
| 97 | + | ||
| 98 | + def test_transform_bias_rescale_qkv(self, device, dtype): | ||
| 99 | + for use_padding in (False, True): | ||
| 100 | + with self.subTest(use_padding=use_padding): | ||
| 101 | + self._test_transform_bias_rescale_qkv_impl( | ||
| 102 | + device, dtype, use_nt=False, use_padding=use_padding | ||
| 103 | + ) | ||
| 104 | + | ||
| 105 | + | ||
| 106 | + | ||
| 107 | + | ||
| 108 | + | ||
| 109 | + | ||
| 110 | + def test_transform_bias_rescale_qkv_nested(self, device, dtype): | ||
| 111 | + for use_padding in (False, True): | ||
| 112 | + with self.subTest(use_padding=use_padding): | ||
| 113 | + self._test_transform_bias_rescale_qkv_impl( | ||
| 114 | + device, dtype, use_nt=True, use_padding=use_padding | ||
| 115 | + ) | ||
| 116 | + | ||
| 117 | + # pylint:disable = huawei-too-many-arguments | ||
| 118 | + def _test_multihead_attention_impl( | ||
| 119 | + self, device, dtype, mode, use_nt, need_weights, average_attn_weights, use_padding=False, pad_all=False | ||
| 120 | + ): | ||
| 121 | + embed_dim = 64 | ||
| 122 | + num_heads = 4 | ||
| 123 | + bs = 16 | ||
| 124 | + sl = 8 | ||
| 125 | + | ||
| 126 | + q = 6 * torch.rand(bs, sl, embed_dim, device=device, dtype=torch.float32) - 3 | ||
| 127 | + if use_padding: | ||
| 128 | + if pad_all: | ||
| 129 | + for q_i in q: | ||
| 130 | + q_i[-1] = torch.zeros_like(q[0][-1], device=device, dtype=torch.float32) | ||
| 131 | + mask = torch.zeros(q.shape[:-1], device=device, dtype=torch.bool) | ||
| 132 | + for mask_i in mask: | ||
| 133 | + mask_i[-1] = True | ||
| 134 | + else: | ||
| 135 | + q[0][-1] = torch.zeros_like(q[0][-1], device=device, dtype=torch.float32) | ||
| 136 | + mask = torch.zeros(q.shape[:-1], device=device, dtype=torch.bool) | ||
| 137 | + mask[0][-1] = True | ||
| 138 | + if mode == "self": | ||
| 139 | + k = q | ||
| 140 | + v = q | ||
| 141 | + elif mode == "encdec": | ||
| 142 | + k = 6 * torch.rand(bs, sl, embed_dim, device=device, dtype=torch.float32) - 3 | ||
| 143 | + v = k | ||
| 144 | + elif mode == "generic": | ||
| 145 | + k = 6 * torch.rand(bs, sl, embed_dim, device=device, dtype=torch.float32) - 3 | ||
| 146 | + v = 6 * torch.rand(bs, sl, embed_dim, device=device, dtype=torch.float32) - 3 | ||
| 147 | + else: | ||
| 148 | + self.fail(f"invalid mode `{mode}`!") | ||
| 149 | + | ||
| 150 | + qkv = torch.nn.Linear(embed_dim, 3 * embed_dim, device=device, dtype=torch.float32) | ||
| 151 | + native_qkv = copy.deepcopy(qkv).to(dtype=dtype) | ||
| 152 | + | ||
| 153 | + proj = torch.nn.Linear(embed_dim, embed_dim, device=device, dtype=torch.float32) | ||
| 154 | + native_proj = copy.deepcopy(proj).to(dtype=dtype) | ||
| 155 | + | ||
| 156 | + pt = torch.nn.MultiheadAttention( | ||
| 157 | + embed_dim, num_heads, batch_first=True, device=device, dtype=torch.float32 | ||
| 158 | + ) | ||
| 159 | + | ||
| 160 | + pt.in_proj_weight = qkv.weight | ||
| 161 | + pt.in_proj_bias = qkv.bias | ||
| 162 | + pt.out_proj.weight = proj.weight | ||
| 163 | + pt.out_proj.bias = proj.bias | ||
| 164 | + | ||
| 165 | + class NativeMHA(torch.nn.Module): | ||
| 166 | + def __init__(self, embed_dim, num_heads, qkv, proj): | ||
| 167 | + super().__init__() | ||
| 168 | + self.qkv = qkv | ||
| 169 | + self.proj = proj | ||
| 170 | + self.embed_dim = embed_dim | ||
| 171 | + self.num_heads = num_heads | ||
| 172 | + | ||
| 173 | + def forward(self, q, k, v, key_padding_mask): | ||
| 174 | + return torch._native_multi_head_attention( | ||
| 175 | + q, | ||
| 176 | + k, | ||
| 177 | + v, | ||
| 178 | + self.embed_dim, | ||
| 179 | + self.num_heads, | ||
| 180 | + self.qkv.weight, | ||
| 181 | + self.qkv.bias, | ||
| 182 | + self.proj.weight, | ||
| 183 | + self.proj.bias, | ||
| 184 | + key_padding_mask, | ||
| 185 | + need_weights=need_weights, | ||
| 186 | + average_attn_weights=average_attn_weights, | ||
| 187 | + mask_type=1, # mask_type = 1 => src_key_padding_mask, mask_type = 0 => src_mask | ||
| 188 | + ) | ||
| 189 | + | ||
| 190 | + npt = NativeMHA( | ||
| 191 | + embed_dim=embed_dim, num_heads=num_heads, qkv=native_qkv, proj=native_proj | ||
| 192 | + ).to(dtype) | ||
| 193 | + | ||
| 194 | + if device == "npu": | ||
| 195 | + pt = pt.npu() | ||
| 196 | + npt = npt.npu() | ||
| 197 | + | ||
| 198 | + ypt, weight_pt = pt( | ||
| 199 | + q, | ||
| 200 | + k, | ||
| 201 | + v, | ||
| 202 | + need_weights=need_weights, | ||
| 203 | + average_attn_weights=average_attn_weights, | ||
| 204 | + key_padding_mask=mask if use_padding else None, | ||
| 205 | + ) | ||
| 206 | + if use_nt: | ||
| 207 | + qs = list(torch.unbind(q)) | ||
| 208 | + if use_padding: | ||
| 209 | + if pad_all: | ||
| 210 | + qs = [x[:-1] for x in qs] | ||
| 211 | + else: | ||
| 212 | + qs[0] = qs[0][:-1] | ||
| 213 | + q = torch.nested.nested_tensor(qs, device=device, dtype=dtype) | ||
| 214 | + if mode == "self": | ||
| 215 | + k = v = q | ||
| 216 | + elif mode == "encdec": | ||
| 217 | + k = torch.nested.nested_tensor(torch.unbind(k), device=device, dtype=dtype) | ||
| 218 | + v = k | ||
| 219 | + else: | ||
| 220 | + k = torch.nested.nested_tensor(torch.unbind(k), device=device, dtype=dtype) | ||
| 221 | + v = torch.nested.nested_tensor(torch.unbind(v), device=device, dtype=dtype) | ||
| 222 | + | ||
| 223 | + native_q = q.to(dtype=dtype) | ||
| 224 | + native_k = k.to(dtype=dtype) | ||
| 225 | + native_v = v.to(dtype=dtype) | ||
| 226 | + | ||
| 227 | + ynpt, weight_npt = npt( | ||
| 228 | + native_q, native_k, native_v, key_padding_mask=mask if use_padding and not use_nt else None | ||
| 229 | + ) | ||
| 230 | + if use_nt: | ||
| 231 | + ynpt = ynpt.to_padded_tensor(0) | ||
| 232 | + if pad_all: | ||
| 233 | + ynpt_final = torch.zeros_like(ypt) | ||
| 234 | + ynpt_final[:, :ynpt.shape[1], :] = ynpt | ||
| 235 | + ynpt = ynpt_final | ||
| 236 | + | ||
| 237 | + def do_pad_all(tensors): | ||
| 238 | + for t in tensors: | ||
| 239 | + for t_i in t: | ||
| 240 | + t_i[-1] = torch.zeros_like(t_i[-1], device=device, dtype=dtype) | ||
| 241 | + | ||
| 242 | + # PyTorch implementation returns non-zero junk in the padding | ||
| 243 | + # locations; overwrite it so that the comparison works out. | ||
| 244 | + if use_padding: | ||
| 245 | + ypt[0][-1] = torch.zeros_like(ypt[0][-1], device=device, dtype=dtype) | ||
| 246 | + ynpt[0][-1] = torch.zeros_like(ynpt[0][-1], device=device, dtype=dtype) | ||
| 247 | + if pad_all: | ||
| 248 | + do_pad_all((ypt, ynpt)) | ||
| 249 | + # Zero the last row of each TxT weight matrix | ||
| 250 | + if need_weights: | ||
| 251 | + if average_attn_weights: | ||
| 252 | + weight_pt[0][-1] = torch.zeros_like(weight_pt[0][-1], device=device, dtype=dtype) | ||
| 253 | + weight_npt[0][-1] = torch.zeros_like(weight_npt[0][-1], device=device, dtype=dtype) | ||
| 254 | + if pad_all: | ||
| 255 | + do_pad_all((weight_pt, weight_npt)) | ||
| 256 | + else: | ||
| 257 | + for nh in range(num_heads): | ||
| 258 | + weight_pt[0][nh][-1] = torch.zeros_like(weight_pt[0][nh][-1], device=device, dtype=dtype) | ||
| 259 | + weight_npt[0][nh][-1] = torch.zeros_like(weight_npt[0][nh][-1], device=device, dtype=dtype) | ||
| 260 | + | ||
| 261 | + if dtype == torch.half: | ||
| 262 | + torch.testing.assert_close(ypt, ynpt.to(torch.float32), atol=1e-3, rtol=1e-3) | ||
| 263 | + else: | ||
| 264 | + # High rtol seems necessary for | ||
| 265 | + # test_native_multihead_attention_cpu_float32 on Windows, | ||
| 266 | + # otherwise 2e-4 would likely be fine. | ||
| 267 | + torch.testing.assert_close(ypt, ynpt, atol=2e-5, rtol=2e-3) | ||
| 268 | + | ||
| 269 | + if need_weights: | ||
| 270 | + torch.testing.assert_close(weight_pt, weight_npt.to(torch.float32), atol=5e-4, rtol=5e-4) | ||
| 271 | + else: | ||
| 272 | + self.assertEqual(weight_pt, weight_npt) | ||
| 273 | + | ||
| 274 | + # NPU currently do not support nested tensor, we set use_nt=False. | ||
| 275 | + # NPU currently do not support calculate with key_padding_mask, we set use_padding=False. | ||
| 276 | + | ||
| 277 | + | ||
| 278 | + | ||
| 279 | + | ||
| 280 | + | ||
| 281 | + | ||
| 282 | + | ||
| 283 | + | ||
| 284 | + | ||
| 285 | + # pylint:disable = huawei-too-many-arguments | ||
| 286 | + def test_native_multihead_self_attention(self, device, dtype, use_nt, | ||
| 287 | + need_weights, average_attn_weights, use_padding, pad_all, fused): | ||
| 288 | + for need_weights in (False, not pad_all): | ||
| 289 | + with self.subTest(use_padding=use_padding, pad_all=pad_all, | ||
| 290 | + use_nt=use_nt, need_weights=need_weights, | ||
| 291 | + average_attn_weights=average_attn_weights): | ||
| 292 | + # NPU do not use sdp_kernel, here we simply call _test_multihead_attention_impl. | ||
| 293 | + if "npu" in device: | ||
| 294 | + self._test_multihead_attention_impl( | ||
| 295 | + device, | ||
| 296 | + dtype, | ||
| 297 | + "self", | ||
| 298 | + use_nt=use_nt, | ||
| 299 | + use_padding=use_padding, | ||
| 300 | + pad_all=pad_all, | ||
| 301 | + need_weights=need_weights, | ||
| 302 | + average_attn_weights=average_attn_weights, | ||
| 303 | + ) | ||
| 304 | + else: | ||
| 305 | + with torch.backends.npu.sdp_kernel( | ||
| 306 | + enable_flash=False, enable_mem_efficient=False | ||
| 307 | + ) if not fused else torch.backends.npu.sdp_kernel( | ||
| 308 | + enable_flash=True, enable_mem_efficient=True | ||
| 309 | + ): | ||
| 310 | + self._test_multihead_attention_impl( | ||
| 311 | + device, | ||
| 312 | + dtype, | ||
| 313 | + "self", | ||
| 314 | + use_nt=use_nt, | ||
| 315 | + use_padding=use_padding, | ||
| 316 | + pad_all=pad_all, | ||
| 317 | + need_weights=need_weights, | ||
| 318 | + average_attn_weights=average_attn_weights, | ||
| 319 | + ) | ||
| 320 | + | ||
| 321 | + | ||
| 322 | + | ||
| 323 | + | ||
| 324 | + | ||
| 325 | + def test_native_multihead_encoder_decoder_attention(self, device, dtype): | ||
| 326 | + self._test_multihead_attention_impl( | ||
| 327 | + device, | ||
| 328 | + dtype, | ||
| 329 | + "encdec", | ||
| 330 | + use_nt=False, | ||
| 331 | + need_weights=False, | ||
| 332 | + average_attn_weights=False, | ||
| 333 | + ) | ||
| 334 | + | ||
| 335 | + | ||
| 336 | + | ||
| 337 | + | ||
| 338 | + | ||
| 339 | + def test_native_multihead_attention(self, device, dtype): | ||
| 340 | + self._test_multihead_attention_impl( | ||
| 341 | + device, | ||
| 342 | + dtype, | ||
| 343 | + "generic", | ||
| 344 | + use_nt=False, | ||
| 345 | + need_weights=False, | ||
| 346 | + average_attn_weights=False, | ||
| 347 | + ) | ||
| 348 | + | ||
| 349 | + | ||
| 350 | +instantiate_device_type_tests(TestMHADeviceType, globals()) | ||
| 351 | + | ||
| 352 | +if __name__ == "__main__": | ||
| 353 | + pass | ||