Initialize a TransformerBlock. Args: layer_id (int): Identifier for the layer. args (ModelArgs): Model configuration parameters. Attributes: n_heads (int): Number of attention heads. dim (int): Dimension size of the model.
(self, layer_id: int, args: ModelArgs)
| 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, |
nothing calls this directly
no test coverage detected