Spaces:
Runtime error
Runtime error
File size: 3,131 Bytes
12edc27 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import math
import torch
from einops import rearrange
from torch import Tensor
# Efficient implementation equivalent to the following:
def scaled_dot_product_attention(query, key, value=None, attn_mask=None, dropout_p=0.0,
is_causal=False, scale=None, enable_gqa=False, return_qk_logits=False) -> torch.Tensor:
L, S = query.size(-2), key.size(-2)
scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
if is_causal:
assert attn_mask is None
temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
attn_bias.to(query.dtype)
if attn_mask is not None:
if attn_mask.dtype == torch.bool:
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
else:
attn_bias = attn_mask + attn_bias
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
attn_weight = query @ key.transpose(-2, -1) * scale_factor
attn_weight += attn_bias
attn_weight = torch.softmax(attn_weight, dim=-1)
attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
if return_qk_logits:
return attn_weight
else:
return attn_weight @ value
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, task_register_embeddings: Tensor = None, use_flash_attention: bool = False) -> Tensor:
q, k = apply_rope(q, k, pe)
if task_register_embeddings is not None:
k_mask_type_embeds, v_mask_type_embeds = task_register_embeddings[0], task_register_embeddings[1]
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
x = rearrange(x, "B H L D -> B L (H D)")
return x
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
assert dim % 2 == 0
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
omega = 1.0 / (theta**scale)
out = torch.einsum("...n,d->...nd", pos, omega)
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
return out.float()
def apply_rope(xq, xk, freqs_cis):
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
if xk is not None:
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
result_q = xq_out.reshape(*xq.shape).type_as(xq)
result_k = xk_out.reshape(*xk.shape).type_as(xk)
return result_q, result_k
else:
result_q = xq_out.reshape(*xq.shape).type_as(xq)
return result_q, None
def apply_learnable_pos_emb(q: Tensor, k: Tensor, embeddings: Tensor) -> Tensor:
if embeddings.ndim == 3:
embeddings = embeddings.unsqueeze(1)
q = q + embeddings
k = k + embeddings
return q, k
|