Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from typing import Tuple | |
| import torch | |
| def rotate_half(x): | |
| x1, x2 = x.chunk(2, dim=-1) | |
| return torch.cat((-x2, x1), dim=-1) | |
| def apply_rotary_pos_emb(x, cos, sin): | |
| cos = cos[:, : x.shape[-2], :] | |
| sin = sin[:, : x.shape[-2], :] | |
| return (x * cos) + (rotate_half(x) * sin) | |
| class RotaryEmbedding(torch.nn.Module): | |
| """ | |
| The rotary position embeddings from RoFormer_ (Su et. al). | |
| A crucial insight from the method is that the query and keys are | |
| transformed by rotation matrices which depend on the relative positions. | |
| Other implementations are available in the Rotary Transformer repo_ and in | |
| GPT-NeoX_, GPT-NeoX was an inspiration | |
| .. _RoFormer: https://arxiv.org/abs/2104.09864 | |
| .. _repo: https://github.com/ZhuiyiTechnology/roformer | |
| .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox | |
| .. warning: Please note that this embedding is not registered on purpose, as it is transformative | |
| (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis | |
| """ | |
| def __init__(self, dim: int, *_, **__): | |
| super().__init__() | |
| # Generate and save the inverse frequency buffer (non trainable) | |
| inv_freq = 1.0 / (10000 ** (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 _update_cos_sin_tables(self, x, seq_dimension=1): | |
| seq_len = x.shape[seq_dimension] | |
| # Reset the tables if the sequence length has changed, | |
| # or if we're on a new device (possibly due to tracing for instance) | |
| if seq_len != self._seq_len_cached or self._cos_cached.device != x.device: | |
| self._seq_len_cached = seq_len | |
| t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq) | |
| freqs = torch.einsum("i,j->ij", t, self.inv_freq) | |
| emb = torch.cat((freqs, freqs), dim=-1).to(x.device) | |
| self._cos_cached = emb.cos()[None, :, :] | |
| self._sin_cached = emb.sin()[None, :, :] | |
| return self._cos_cached, self._sin_cached | |
| def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k, seq_dimension=-2) | |
| return ( | |
| apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), | |
| apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), | |
| ) | |