Implementation of RotaryEmbedding from GPT-NeoX. This implementation is design to operate on queries and keys that are compatible with [batch_size, n_heads_per_partition, seq_len, head_dim] (e.g. MinGPTAttention format).
| 25 | |
| 26 | |
| 27 | class RotaryEmbedding(nn.Module): |
| 28 | """Implementation of RotaryEmbedding from GPT-NeoX. |
| 29 | This implementation is design to operate on queries and keys that are compatible with |
| 30 | [batch_size, n_heads_per_partition, seq_len, head_dim] (e.g. MinGPTAttention format). |
| 31 | """ |
| 32 | |
| 33 | def __init__( |
| 34 | self, |
| 35 | head_dim: int, |
| 36 | base=10000, |
| 37 | ): |
| 38 | super().__init__() |
| 39 | inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| 40 | self.register_buffer("inv_freq", inv_freq, persistent=False) |
| 41 | self.head_dim = head_dim |
| 42 | self.seq_len_cached = None |
| 43 | self.batch_size_cached = None |
| 44 | self.cos_cached: torch.Tensor | None = None |
| 45 | self.sin_cached: torch.Tensor | None = None |
| 46 | |
| 47 | def cos_sin( |
| 48 | self, |
| 49 | seq_len: int, |
| 50 | device="cuda", |
| 51 | dtype=torch.bfloat16, |
| 52 | ) -> torch.Tensor: |
| 53 | if seq_len != self.seq_len_cached: |
| 54 | self.seq_len_cached = seq_len |
| 55 | t = torch.arange(seq_len, device=device).type_as(self.inv_freq) |
| 56 | freqs = torch.einsum("i,j->ij", t, self.inv_freq) |
| 57 | emb = torch.cat((freqs, freqs), dim=-1).to(device) |
| 58 | |
| 59 | if dtype in [torch.float16, torch.bfloat16]: |
| 60 | emb = emb.float() |
| 61 | |
| 62 | self.cos_cached = emb.cos()[None, :, :] |
| 63 | self.sin_cached = emb.sin()[None, :, :] |
| 64 | |
| 65 | self.cos_cached = self.cos_cached.type(dtype) |
| 66 | self.sin_cached = self.sin_cached.type(dtype) |
| 67 | |
| 68 | return self.cos_cached, self.sin_cached |
| 69 | |
| 70 | def forward(self, _q, _k): |
| 71 | batch, seq_len, num_heads, head_dim = _q.shape |
| 72 | q = _q.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim) |
| 73 | k = _k.permute(0, 2, 1, 3).contiguous().reshape(-1, seq_len, head_dim) |
| 74 | cos, sin = self.cos_sin(seq_len, q.device, q.dtype) |
| 75 | return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) |
| 76 | |
| 77 | |
| 78 | class FalconAttentionFused(nn.Module): |