| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | import torch.nn.functional as F |
| 5 | from torch import nn |
| 6 | |
| 7 | from vllm.distributed import ( |
| 8 | get_tensor_model_parallel_rank, |
| 9 | get_tensor_model_parallel_world_size, |
| 10 | ) |
| 11 | |
| 12 | from .base import RotaryEmbedding |
| 13 | from .common import rotate_neox |
| 14 | |
| 15 | |
| 16 | class FourierRotaryEmbedding(RotaryEmbedding): |
| 17 | def __init__( |
| 18 | self, |
| 19 | head_size: int, |
| 20 | rotary_dim: int, |
| 21 | max_position_embeddings: int, |
| 22 | base: float, |
| 23 | is_neox_style: bool, |
| 24 | dtype: torch.dtype, |
| 25 | init_cache: bool, |
| 26 | # extra parameters for FoPE |
| 27 | num_key_value_heads: int, |
| 28 | num_inv_freq: int, |
| 29 | fope_sep_head: bool, |
| 30 | fope_init_factor: float, |
| 31 | ): |
| 32 | # fope related parameters |
| 33 | self.num_key_value_heads = num_key_value_heads |
| 34 | self.num_inv_freq = num_inv_freq |
| 35 | self.fope_sep_head = fope_sep_head |
| 36 | self.fope_init_factor = fope_init_factor |
| 37 | |
| 38 | super().__init__( |
| 39 | head_size=head_size, |
| 40 | rotary_dim=rotary_dim, |
| 41 | max_position_embeddings=max_position_embeddings, |
| 42 | base=base, |
| 43 | is_neox_style=is_neox_style, |
| 44 | dtype=dtype, |
| 45 | init_cache=init_cache, |
| 46 | ) |
| 47 | |
| 48 | # setup buffers and parameters |
| 49 | self.inv_freq: torch.Tensor |
| 50 | self.register_buffer( |
| 51 | "inv_freq", self._compute_inv_freq(self.base), persistent=False |
| 52 | ) |
| 53 | |
| 54 | self.input_dim = self.inv_freq.shape[-1] |
| 55 | self.output_dim = self.inv_freq.shape[-1] |
| 56 | self.cos_coef = nn.Parameter( |
| 57 | torch.empty(num_key_value_heads, self.input_dim, self.output_dim), |
| 58 | requires_grad=False, |
| 59 | ) |
| 60 | self.sin_coef = nn.Parameter( |
| 61 | torch.empty(num_key_value_heads, self.input_dim, self.output_dim), |
| 62 | requires_grad=False, |
| 63 | ) |
| 64 | self.sin_coef.weight_loader = self.weight_loader |
| 65 | self.cos_coef.weight_loader = self.weight_loader |
| 66 | |
| 67 | self.cos_sin_cache: torch.Tensor |
| 68 | cache = self._compute_cos_sin_cache().to(dtype) |
| 69 | self.register_buffer("cos_sin_cache", cache, persistent=False) |
| 70 | |
| 71 | # update cache in the first forward, where sin/cos_coef weights are ready |
| 72 | self.update_cache = True |
| 73 | |
| 74 | def _compute_inv_freq(self, base: float) -> torch.Tensor: |
| 75 | """Compute the inverse frequency.""" |
| 76 | inv_freq = 1.0 / ( |
| 77 | base |
| 78 | ** ( |
| 79 | torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim |
| 80 | ) |
| 81 | ) |
| 82 | |
| 83 | inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) |
| 84 | if self.num_inv_freq is not None: |
| 85 | inv_freq_idx_selected[self.num_inv_freq :] = False |
| 86 | else: |
| 87 | inv_freq_idx_selected = inv_freq > ( |
| 88 | 2.0 * torch.pi / self.max_position_embeddings |
| 89 | ) |
| 90 | |
| 91 | inv_freq = inv_freq[inv_freq_idx_selected] |
| 92 | return inv_freq |
| 93 | |
| 94 | def _compute_cos_sin_cache(self) -> torch.Tensor: |
| 95 | """Compute the cos and sin cache.""" |
| 96 | device = self.inv_freq.device |
| 97 | t = torch.arange(self.max_position_embeddings, dtype=torch.float, device=device) |
| 98 | |
| 99 | freqs = torch.einsum("j,i -> ji", t, self.inv_freq) |
| 100 | if self.fope_sep_head: |
| 101 | pos_cos = freqs.cos().unsqueeze(0).expand(self.num_key_value_heads, -1, -1) |
| 102 | pos_sin = freqs.sin().unsqueeze(0).expand(self.num_key_value_heads, -1, -1) |
| 103 | else: |
| 104 | pos_cos = freqs.cos() |
| 105 | pos_sin = freqs.sin() |
| 106 | |
| 107 | if self.fope_sep_head: |
| 108 | sin = torch.einsum("htD, hDd -> thd", pos_sin, self.sin_coef.float()) |
| 109 | cos = torch.einsum("htD, hDd -> thd", pos_cos, self.cos_coef.float()) |
| 110 | else: |
| 111 | sin = torch.einsum("tD, Dd -> td", pos_sin, self.sin_coef.float()) |
| 112 | cos = torch.einsum("tD, Dd -> td", pos_cos, self.cos_coef.float()) |
| 113 | |
| 114 | sin = F.pad( |
| 115 | input=sin, |
| 116 | pad=(0, self.head_size // 2 - sin.size(-1)), |
| 117 | mode="constant", |
| 118 | value=1, |
| 119 | ) |
| 120 | cos = F.pad( |
| 121 | input=cos, |
| 122 | pad=(0, self.head_size // 2 - cos.size(-1)), |
| 123 | mode="constant", |
| 124 | value=1, |
| 125 | ) |
| 126 | |
| 127 | sin = torch.cat((sin, sin), dim=-1) |
| 128 | cos = torch.cat((cos, cos), dim=-1) |
| 129 | |
| 130 | # cache: (max_position_embeddings, num_kv_heads, kv_size * 2) |
| 131 | cache = torch.cat((cos, sin), dim=-1) |
| 132 | return cache |
| 133 | |
| 134 | def forward_native( |
| 135 | self, |
| 136 | positions: torch.Tensor, |
| 137 | query: torch.Tensor, |
| 138 | key: torch.Tensor | None = None, |
| 139 | offsets: torch.Tensor | None = None, |
| 140 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 141 | # update cos/sin cache in the first forward |
| 142 | if self.update_cache: |
| 143 | cache = self._compute_cos_sin_cache().to(self.dtype) |
| 144 | self.cos_sin_cache.copy_(cache) |
| 145 | self.update_cache = False |
| 146 | |
| 147 | positions = positions.flatten() |
| 148 | cos_sin = self.cos_sin_cache.index_select(0, positions) |
| 149 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 150 | |
| 151 | # apply rotary embedding |
| 152 | # query: (seq_len, num_heads, head_size) |
| 153 | # key: (seq_len, num_kv_heads, head_size) |
| 154 | query = query.unflatten(-1, (-1, self.head_size)) |
| 155 | assert key is not None, "Key tensor is required for FoPE." |
| 156 | key = key.unflatten(-1, (-1, self.head_size)) |
| 157 | |
| 158 | assert query.dim() == key.dim() == 3, ( |
| 159 | "Expected query key (seq_len, heads, head_dim)" |
| 160 | ) |
| 161 | assert cos.dim() <= 3 and sin.dim() <= 3 |
| 162 | |
| 163 | need_reshape = False |
| 164 | if cos.dim() == 3: |
| 165 | # for fope |
| 166 | need_reshape = True |
| 167 | query_shape = query.shape |
| 168 | key_shape = key.shape |
| 169 | cos = cos.flatten(0, 1) |
| 170 | sin = sin.flatten(0, 1) |
| 171 | seq_len = cos.size(0) |
| 172 | query = query.view(seq_len, -1, query.size(-1)) |
| 173 | key = key.view(seq_len, -1, key.size(-1)) |
| 174 | |
| 175 | # native implementation of apply rope for neox style |
| 176 | cos = cos.unsqueeze(1) |
| 177 | sin = sin.unsqueeze(1) |
| 178 | query = (query * cos) + (rotate_neox(query) * sin) |
| 179 | key = (key * cos) + (rotate_neox(key) * sin) |
| 180 | |
| 181 | if need_reshape: |
| 182 | query = query.view(query_shape) |
| 183 | key = key.view(key_shape) |
| 184 | |
| 185 | return query, key |
| 186 | |
| 187 | def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): |
| 188 | """Load fope weights.""" |
| 189 | world_size = get_tensor_model_parallel_world_size() |
| 190 | rank = get_tensor_model_parallel_rank() |
| 191 | num_key_value_heads = loaded_weight.size(0) |
| 192 | |
| 193 | if num_key_value_heads < world_size: |
| 194 | n_replicate = world_size // num_key_value_heads |
| 195 | world_size = num_key_value_heads |
| 196 | rank = rank // n_replicate |
| 197 | |
| 198 | loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] |
| 199 | param.data.copy_(loaded_weight) |