File size: 4,721 Bytes
7f316fe | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
# Part of this implementation is adapted from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py#L20
# which is released under BSD-3 license
# Part of this implementation is adapted from https://github.com/louaaron/Score-Entropy-Discrete-Diffusion
# which is released under MIT license
from typing import Tuple, Optional
import torch
from einops import repeat
from torch import Tensor
class Rotary(torch.nn.Module):
"""
From: https://github.com/louaaron/Score-Entropy-Discrete-Diffusion
"""
def __init__(self, dim: int, base: int = 10_000):
super().__init__()
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
self.register_buffer("inv_freq", inv_freq)
self.seq_len_cached = None
self.cos_cached = None
self.sin_cached = None
def positions_like(self, lengths: torch.Tensor) -> torch.Tensor:
# returns (T,) positions 0..len_i-1 concatenated across sequences
return torch.cat([torch.arange(int(L), device=lengths.device) for L in lengths], dim=0)
def forward(self, x: Tensor, seq_dim: int = 1) -> Tuple[Tensor, Tensor]:
seq_len = x.shape[seq_dim]
if seq_len != self.seq_len_cached:
self.seq_len_cached = seq_len
t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq)
freqs = torch.einsum("i,j->ij", t, self.inv_freq.clone())
emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
# dims are: batch, seq_len, qkv, head, dim
self.cos_cached = emb.cos()[None, :, None, None, :].repeat(1, 1, 3, 1, 1)
self.sin_cached = emb.sin()[None, :, None, None, :].repeat(1, 1, 3, 1, 1)
# This makes the transformation on v an identity.
self.cos_cached[:, :, 2, :, :].fill_(1.0)
self.sin_cached[:, :, 2, :, :].fill_(0.0)
return self.cos_cached, self.sin_cached
def rotate_half(x: Tensor) -> Tensor:
x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_emb_torch(x, cos, sin, interleaved: bool = False, rotary_dim: Optional[int] = None):
# cos/sin cached as in your Rotary.forward
cos = cos[0, :, 0, 0, : cos.shape[-1] // 2].to(dtype=x.dtype, device=x.device) # (S, d/2)
sin = sin[0, :, 0, 0, : sin.shape[-1] // 2].to(dtype=x.dtype, device=x.device) # (S, d/2)
# Decide how many head dims get RoPE
if rotary_dim is None:
rotary_dim = cos.shape[-1] * 2 # default: full cached rotary span
rotary_dim = min(rotary_dim, x.shape[-1])
assert rotary_dim % 2 == 0, "rotary_dim must be even"
# Broadcast to (..., S, 1, rotary_dim)
cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
x_rot = x[..., :rotary_dim]
x_tail = x[..., rotary_dim:]
x_rot = x_rot * cos + rotate_half(x_rot) * sin
return torch.cat([x_rot, x_tail], dim=-1)
def apply_rotary_emb_ragged(
q: torch.Tensor, k: torch.Tensor,
cos: torch.Tensor, sin: torch.Tensor,
positions: torch.Tensor, head_dim: Optional[int] = None,
interleaved: bool = False, inplace: bool = True
):
# q,k: (T, H, Dh), positions: (T,)
cos_half = cos[0, :, 0, 0, : cos.shape[-1] // 2].to(dtype=q.dtype, device=q.device) # (S, d/2)
sin_half = sin[0, :, 0, 0, : sin.shape[-1] // 2].to(dtype=q.dtype, device=q.device)
assert positions.max().item() < cos_half.size(0), "position index exceeds cached length"
cos_sel = cos_half.index_select(0, positions) # (T, d/2)
sin_sel = sin_half.index_select(0, positions) # (T, d/2)
cos_sel = repeat(cos_sel, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") # (T,1,d)
sin_sel = repeat(sin_sel, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
ro_dim = head_dim or cos_sel.shape[-1]
ro_dim = min(ro_dim, q.shape[-1])
assert ro_dim % 2 == 0, "rotary_dim/head_dim must be even"
if not inplace:
q = q.clone()
k = k.clone()
def rot(x):
x1, x2 = x[..., : ro_dim // 2], x[..., ro_dim // 2: ro_dim]
return torch.cat([-x2, x1], dim=-1)
q_head = q[..., :ro_dim]
k_head = k[..., :ro_dim]
q[..., :ro_dim] = q_head * cos_sel + rot(q_head) * sin_sel
k[..., :ro_dim] = k_head * cos_sel + rot(k_head) * sin_sel
return q, k
|