| 355 | |
| 356 | |
| 357 | class TransformerBlock(nn.Module): |
| 358 | def __init__(self, layer_id: int, args: ModelArgs): |
| 359 | """ |
| 360 | Initialize a TransformerBlock. |
| 361 | |
| 362 | Args: |
| 363 | layer_id (int): Identifier for the layer. |
| 364 | args (ModelArgs): Model configuration parameters. |
| 365 | |
| 366 | Attributes: |
| 367 | n_heads (int): Number of attention heads. |
| 368 | dim (int): Dimension size of the model. |
| 369 | head_dim (int): Dimension size of each attention head. |
| 370 | attention (Attention): Attention module. |
| 371 | feed_forward (FeedForward): FeedForward module. |
| 372 | layer_id (int): Identifier for the layer. |
| 373 | attention_norm (RMSNorm): Layer normalization for attention output. |
| 374 | ffn_norm (RMSNorm): Layer normalization for feedforward output. |
| 375 | |
| 376 | """ |
| 377 | super().__init__() |
| 378 | self.n_heads = args.n_heads |
| 379 | self.dim = args.dim |
| 380 | self.head_dim = args.dim // args.n_heads |
| 381 | self.attention = Attention(args) |
| 382 | self.feed_forward = FeedForward( |
| 383 | dim=args.dim, |
| 384 | hidden_dim=4 * args.dim, |
| 385 | multiple_of=args.multiple_of, |
| 386 | ffn_dim_multiplier=args.ffn_dim_multiplier, |
| 387 | ) |
| 388 | self.layer_id = layer_id |
| 389 | self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps) |
| 390 | self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps) |
| 391 | |
| 392 | def forward( |
| 393 | self, |
| 394 | x: torch.Tensor, |
| 395 | start_pos: int, |
| 396 | freqs_cis: torch.Tensor, |
| 397 | mask: Optional[torch.Tensor], |
| 398 | ): |
| 399 | """ |
| 400 | Perform a forward pass through the TransformerBlock. |
| 401 | |
| 402 | Args: |
| 403 | x (torch.Tensor): Input tensor. |
| 404 | start_pos (int): Starting position for attention caching. |
| 405 | freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies. |
| 406 | mask (torch.Tensor, optional): Masking tensor for attention. Defaults to None. |
| 407 | |
| 408 | Returns: |
| 409 | torch.Tensor: Output tensor after applying attention and feedforward layers. |
| 410 | |
| 411 | """ |
| 412 | h = x + self.attention.forward( |
| 413 | self.attention_norm(x), start_pos, freqs_cis, mask |
| 414 | ) |