A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
| 419 | |
| 420 | |
| 421 | class QKVAttentionLegacy(nn.Module): |
| 422 | """ |
| 423 | A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping |
| 424 | """ |
| 425 | |
| 426 | def __init__(self, n_heads): |
| 427 | super().__init__() |
| 428 | self.n_heads = n_heads |
| 429 | |
| 430 | def forward(self, qkv): |
| 431 | """ |
| 432 | Apply QKV attention. |
| 433 | :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs. |
| 434 | :return: an [N x (H * C) x T] tensor after attention. |
| 435 | """ |
| 436 | bs, width, length = qkv.shape |
| 437 | assert width % (3 * self.n_heads) == 0 |
| 438 | ch = width // (3 * self.n_heads) |
| 439 | q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1) |
| 440 | scale = 1 / math.sqrt(math.sqrt(ch)) |
| 441 | weight = th.einsum("bct,bcs->bts", q * scale, k * scale) # More stable with f16 than dividing afterwards |
| 442 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) |
| 443 | a = th.einsum("bts,bcs->bct", weight, v) |
| 444 | return a.reshape(bs, -1, length) |
| 445 | |
| 446 | @staticmethod |
| 447 | def count_flops(model, _x, y): |
| 448 | return count_flops_attn(model, _x, y) |
| 449 | |
| 450 | |
| 451 | class QKVAttention(nn.Module): |