| 158 | self.register_buffer("mask", torch.empty((1, 1, 0, 0), dtype=bool)) |
| 159 | |
| 160 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 161 | # batch size, sequence length, embedding dimensionality |
| 162 | B, T, C = x.shape |
| 163 | if C != self.embedding_dim: |
| 164 | raise ValueError( |
| 165 | f"Expected input shape (..., {self.embedding_dim}, got {x.shape})" |
| 166 | ) |
| 167 | |
| 168 | # calculate and separate out query, key, values for all heads |
| 169 | # each has shape (B, T, C) |
| 170 | q, k, v = self.attention(x).split(self.embedding_dim, dim=2) |
| 171 | |
| 172 | # separate out head index and move it up next to the batch dimension |
| 173 | # final shape (B, num_heads, T, head_size), where C = num_heads * head_size |
| 174 | nh = self.num_heads |
| 175 | hs = C // nh |
| 176 | k = k.view(B, T, nh, hs).transpose(1, 2) |
| 177 | q = q.view(B, T, nh, hs).transpose(1, 2) |
| 178 | v = v.view(B, T, nh, hs).transpose(1, 2) |
| 179 | |
| 180 | # self-attention: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) |
| 181 | if self.uses_flash: |
| 182 | # efficient attention using Flash Attention CUDA kernels |
| 183 | dropout_p = self.dropout if self.training else 0 |
| 184 | y = F.scaled_dot_product_attention( |
| 185 | q, k, v, attn_mask=None, dropout_p=dropout_p, is_causal=self.causal |
| 186 | ) |
| 187 | else: |
| 188 | # manual implementation of attention |
| 189 | att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(hs)) |
| 190 | |
| 191 | if self.causal: |
| 192 | # cache the causal mask, if we're using one |
| 193 | if self.mask.shape[2] < T: |
| 194 | self.mask = torch.tril(torch.ones(T, T)).view(1, 1, T, T) == 0 |
| 195 | att = att.masked_fill(self.bias[:, :, :T, :T], float("-inf")) |
| 196 | |
| 197 | att = F.softmax(att, dim=-1) |
| 198 | att = self.attention_dropout(att) |
| 199 | # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) |
| 200 | y = att @ v |
| 201 | |
| 202 | # re-assemble all head outputs side by side |
| 203 | y = y.transpose(1, 2).contiguous().view(B, T, C) |
| 204 | |
| 205 | # output projection |
| 206 | y = self.residual_dropout(self.projection(y)) |
| 207 | return y |
| 208 | |
| 209 | |
| 210 | class TransformerBlock(nn.Module): |