"""rotary_positional_embedding.py - Rotary Positional Embedding code from github.com/lucidrains/rotary-embedding-torch MIT License """ from math import pi, log import torch from torch import nn, einsum from einops import rearrange, repeat def exists(val): return val is not None def broadcat(tensors, dim=-1): num_tensors = len(tensors) shape_lens = set(list(map(lambda t: len(t.shape), tensors))) assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions' shape_len = list(shape_lens)[0] dim = (dim + shape_len) if dim < 0 else dim dims = list(zip(*map(lambda t: list(t.shape), tensors))) expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim] assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims) ]), 'invalid dimensions for broadcastable concatentation' max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims)) expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims)) expanded_dims.insert(dim, (dim, dims[dim])) expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims))) tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes))) return torch.cat(tensors, dim=dim) # rotary embedding helper functions def rotate_half(x): x = rearrange(x, '... (d r) -> ... d r', r=2) x1, x2 = x.unbind(dim=-1) x = torch.stack((-x2, x1), dim=-1) return rearrange(x, '... d r -> ... (d r)') def apply_rotary_emb(freqs, t, start_index=0, scale=1.): rot_dim, seq_len = freqs.shape[-1], t.shape[-2] freqs = freqs[-seq_len:, :] freqs = freqs.to(t) end_index = start_index + rot_dim assert rot_dim <= t.shape[ -1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}' t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:] t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale) return torch.cat((t_left, t, t_right), dim=-1) # learned rotation helpers def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None): if exists(freq_ranges): rotations = einsum('..., f -> ... f', rotations, freq_ranges) rotations = rearrange(rotations, '... r f -> ... (r f)') rotations = repeat(rotations, '... n -> ... (n r)', r=2) return apply_rotary_emb(rotations, t, start_index=start_index) # classes class RotaryEmbedding(nn.Module): def __init__(self, dim, custom_freqs=None, freqs_for='lang', theta=10000, max_freq=10, num_freqs=1, learned_freq=False, use_xpos=False, xpos_scale_base=512, interpolate_factor=1., theta_rescale_factor=1.): super().__init__() # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning # has some connection to NTK literature # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ theta *= theta_rescale_factor**(dim / (dim - 2)) if exists(custom_freqs): freqs = custom_freqs elif freqs_for == 'lang': freqs = 1. / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) elif freqs_for == 'pixel': freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi elif freqs_for == 'constant': freqs = torch.ones(num_freqs).float() else: raise ValueError(f'unknown modality {freqs_for}') self.cache = dict() self.cache_scale = dict() self.freqs = nn.Parameter(freqs, requires_grad=learned_freq) # interpolation factors assert interpolate_factor >= 1. self.interpolate_factor = interpolate_factor # xpos self.use_xpos = use_xpos if not use_xpos: self.register_buffer('scale', None) return scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) self.scale_base = xpos_scale_base self.register_buffer('scale', scale) def get_seq_pos(self, seq_len, device, dtype, offset=0): return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor def rotate_queries_or_keys(self, t, seq_dim=-2, offset=0, freq_seq_len=None): assert not self.use_xpos, 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings' device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim] if exists(freq_seq_len): assert freq_seq_len >= seq_len seq_len = freq_seq_len freqs = self.forward( lambda: self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset), cache_key=f'freqs:{seq_len}|offset:{offset}') return apply_rotary_emb(freqs, t) def rotate_queries_with_cached_keys(self, q, k, seq_dim=-2): q_len, k_len = q.shape[seq_dim], k.shape[seq_dim] assert q_len <= k_len q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, freq_seq_len=k_len) k = self.rotate_queries_or_keys(k, seq_dim=seq_dim) return q, k def rotate_queries_and_keys(self, q, k, seq_dim=-2): assert self.use_xpos device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim] seq = self.get_seq_pos(seq_len, dtype=dtype, device=device) freqs = self.forward(lambda: seq, cache_key=f'freqs:{seq_len}') scale = self.get_scale(lambda: seq, cache_key=f'scale:{seq_len}').to(dtype) rotated_q = apply_rotary_emb(freqs, q, scale=scale) rotated_k = apply_rotary_emb(freqs, k, scale=scale**-1) return rotated_q, rotated_k def get_scale(self, t, cache_key=None): assert self.use_xpos if exists(cache_key) and cache_key in self.cache: return self.cache[cache_key] if callable(t): t = t() scale = 1. if self.use_xpos: power = (t - len(t) // 2) / self.scale_base scale = self.scale**rearrange(power, 'n -> n 1') scale = torch.cat((scale, scale), dim=-1) if exists(cache_key): self.cache[cache_key] = scale return scale def forward(self, t, cache_key=None): if exists(cache_key) and cache_key in self.cache: return self.cache[cache_key] if callable(t): t = t() freqs = self.freqs freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs) freqs = repeat(freqs, '... n -> ... (n r)', r=2) if exists(cache_key): self.cache[cache_key] = freqs return freqs