| 243 | |
| 244 | |
| 245 | class Block(nn.Module): |
| 246 | def __init__(self, config: Config) -> None: |
| 247 | super().__init__() |
| 248 | self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps) |
| 249 | self.attn = CausalSelfAttention(config) |
| 250 | if not config.shared_attention_norm: |
| 251 | self.norm_2 = config.norm_class(config.n_embd, eps=config.norm_eps) |
| 252 | self.mlp = config.mlp_class(config) |
| 253 | self.config = config |
| 254 | |
| 255 | def forward( |
| 256 | self, |
| 257 | x: torch.Tensor, |
| 258 | rope: RoPECache, |
| 259 | max_seq_length: int, |
| 260 | mask: Optional[torch.Tensor] = None, |
| 261 | input_pos: Optional[torch.Tensor] = None, |
| 262 | kv_cache: Optional[KVCache] = None, |
| 263 | ) -> Tuple[torch.Tensor, Optional[KVCache]]: |
| 264 | n_1 = self.norm_1(x) |
| 265 | h, new_kv_cache = self.attn( |
| 266 | n_1, rope, max_seq_length, mask, input_pos, kv_cache |
| 267 | ) |
| 268 | if self.config.parallel_residual: |
| 269 | n_2 = n_1 if self.config.shared_attention_norm else self.norm_2(x) |
| 270 | x = x + h + self.mlp(n_2) |
| 271 | else: |
| 272 | if self.config.shared_attention_norm: |
| 273 | raise NotImplementedError( |
| 274 | 'No checkpoint amongst the ones we support uses this configuration' |
| 275 | ' (non-parallel residual and shared attention norm).' |
| 276 | ) |
| 277 | |
| 278 | x = x + h |
| 279 | x = x + self.mlp(self.norm_2(x)) |
| 280 | return x, new_kv_cache |
| 281 | |
| 282 | |
| 283 | class CausalSelfAttention(nn.Module): |