File size: 3,518 Bytes
cf5d356 | 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 | # 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 __future__ import annotations
import torch
from torch import Tensor, nn
def rotate_half_dimension(inputs: Tensor) -> Tensor:
"""Rotate pairs of feature halves by 90 degrees."""
if inputs.shape[-1] % 2:
raise ValueError("The final input dimension must be even.")
first_half, second_half = inputs.chunk(2, dim=-1)
return torch.cat((-second_half, first_half), dim=-1)
def apply_rotary_embedding(
inputs: Tensor,
cosine: Tensor,
sine: Tensor,
) -> Tensor:
"""Apply precomputed rotary tables to query or key tensors."""
sequence_length = inputs.shape[-2]
cosine = cosine[:, :sequence_length, :]
sine = sine[:, :sequence_length, :]
return inputs * cosine + rotate_half_dimension(inputs) * sine
class RotaryEmbedding(nn.Module):
"""Rotary position embedding used for attention queries and keys."""
def __init__(self, embedding_dimension: int) -> None:
super().__init__()
if embedding_dimension <= 0 or embedding_dimension % 2:
raise ValueError("embedding_dimension must be a positive even integer.")
inverse_frequencies = 1.0 / (
10000
** (
torch.arange(0, embedding_dimension, 2).float()
/ embedding_dimension
)
)
self.register_buffer("inv_freq", inverse_frequencies)
self._cached_sequence_length = 0
self._cached_cosine: Tensor | None = None
self._cached_sine: Tensor | None = None
def _get_cosine_sine_tables(
self,
inputs: Tensor,
sequence_dimension: int = 1,
) -> tuple[Tensor, Tensor]:
sequence_length = inputs.shape[sequence_dimension]
if (
sequence_length != self._cached_sequence_length
or self._cached_cosine is None
or self._cached_sine is None
or self._cached_cosine.device != inputs.device
or self._cached_cosine.dtype != inputs.dtype
):
self._cached_sequence_length = sequence_length
positions = torch.arange(
sequence_length,
device=inputs.device,
dtype=self.inv_freq.dtype,
)
frequencies = torch.einsum(
"i,j->ij", positions, self.inv_freq
)
embedding = torch.cat(
(frequencies, frequencies), dim=-1
).to(dtype=inputs.dtype)
self._cached_cosine = embedding.cos()[None, :, :]
self._cached_sine = embedding.sin()[None, :, :]
return self._cached_cosine, self._cached_sine
def forward(
self,
query: Tensor,
key: Tensor,
) -> tuple[Tensor, Tensor]:
"""Rotate query and key tensors with shape ``[B, H, L, D]``."""
if query.shape != key.shape:
raise ValueError("query and key must have identical shapes.")
self._cached_cosine, self._cached_sine = (
self._get_cosine_sine_tables(
key,
sequence_dimension=-2,
)
)
return (
apply_rotary_embedding(
query, self._cached_cosine, self._cached_sine
),
apply_rotary_embedding(
key, self._cached_cosine, self._cached_sine
),
)
|