| 20 | |
| 21 | |
| 22 | class QuantLlamaRotaryEmbedding(nn.Module): |
| 23 | def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): |
| 24 | super().__init__() |
| 25 | |
| 26 | self.dim = dim |
| 27 | self.max_position_embeddings = max_position_embeddings |
| 28 | self.base = base |
| 29 | inv_freq = 1.0 / ( |
| 30 | self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) |
| 31 | ) |
| 32 | self.register_buffer("inv_freq", inv_freq) |
| 33 | # Build here to make `torch.jit.trace` work. |
| 34 | self._set_cos_sin_cache( |
| 35 | seq_len=max_position_embeddings, |
| 36 | device=self.inv_freq.device, |
| 37 | dtype=torch.get_default_dtype(), |
| 38 | ) |
| 39 | |
| 40 | def _set_cos_sin_cache(self, seq_len, device, dtype): |
| 41 | self.max_seq_len_cached = seq_len |
| 42 | t = torch.arange( |
| 43 | self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype |
| 44 | ) |
| 45 | |
| 46 | freqs = torch.einsum("i,j->ij", t, self.inv_freq) |
| 47 | # Different from paper, but it uses a different permutation in order to obtain the same calculation |
| 48 | emb = torch.cat((freqs, freqs), dim=-1) |
| 49 | |
| 50 | cos = freqs.cos() |
| 51 | sin = freqs.sin() |
| 52 | cache = torch.cat((cos, sin), dim=-1) |
| 53 | |
| 54 | # self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False) |
| 55 | # self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False) |
| 56 | self.register_buffer("cos_sin_cache", cache.half(), persistent=False) |
| 57 | |
| 58 | def forward( |
| 59 | self, |
| 60 | query: torch.Tensor, |
| 61 | key: torch.Tensor, |
| 62 | positions: torch.Tensor, |
| 63 | ): |
| 64 | # Apply rotary embedding to the query and key before passing them |
| 65 | # to the attention op. |
| 66 | # print(positions.shape, query.shape, key.shape, self.cos_sin_cache.shape) |
| 67 | query = query.contiguous() |
| 68 | key = key.contiguous() |
| 69 | awq_inference_engine.rotary_embedding_neox( |
| 70 | positions, |
| 71 | query, |
| 72 | key, |
| 73 | self.dim, |
| 74 | self.cos_sin_cache, |
| 75 | ) |
| 76 | return query, key |
| 77 | |
| 78 | |
| 79 | class QuantLlamaAttention(nn.Module): |