Spaces:
Running
on
Zero
Running
on
Zero
| 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 | |