TransformerBlock Module Args: layer_id (int): Identifier for the layer. model_args (ModelArgs): Model configuration arguments. Attributes: n_heads (int): Number of attention heads. dim (int): Dimension size of the model. head_dim (int): Dimensio
| 273 | |
| 274 | |
| 275 | class TransformerBlock(nn.Module): |
| 276 | """ |
| 277 | TransformerBlock Module |
| 278 | |
| 279 | Args: |
| 280 | layer_id (int): Identifier for the layer. |
| 281 | model_args (ModelArgs): Model configuration arguments. |
| 282 | |
| 283 | Attributes: |
| 284 | n_heads (int): Number of attention heads. |
| 285 | dim (int): Dimension size of the model. |
| 286 | head_dim (int): Dimension size of each attention head. |
| 287 | attention (Attention): Attention module. |
| 288 | feed_forward (FeedForward): FeedForward module. |
| 289 | layer_id (int): Identifier for the layer. |
| 290 | attention_norm (RMSNorm): Layer normalization for attention output. |
| 291 | ffn_norm (RMSNorm): Layer normalization for feedforward output. |
| 292 | |
| 293 | """ |
| 294 | |
| 295 | def __init__(self, layer_id: int, model_args: ModelArgs): |
| 296 | super().__init__() |
| 297 | self.n_heads = model_args.n_heads |
| 298 | self.dim = model_args.dim |
| 299 | self.attention = Attention(model_args) |
| 300 | self.feed_forward = FeedForward( |
| 301 | dim=model_args.dim, |
| 302 | hidden_dim=4 * model_args.dim, |
| 303 | multiple_of=model_args.multiple_of, |
| 304 | ffn_dim_multiplier=model_args.ffn_dim_multiplier, |
| 305 | ) |
| 306 | self.layer_id = layer_id |
| 307 | self.num_layers = model_args.n_layers |
| 308 | |
| 309 | self.attention_norm = RMSNorm( |
| 310 | dim=model_args.dim, eps=model_args.norm_eps |
| 311 | ) |
| 312 | self.ffn_norm = RMSNorm( |
| 313 | dim=model_args.dim, eps=model_args.norm_eps |
| 314 | ) |
| 315 | |
| 316 | if model_args.depth_init: |
| 317 | self.weight_init_std = 0.02 / (2 * (self.layer_id + 1)) ** 0.5 |
| 318 | else: |
| 319 | self.weight_init_std = 0.02 / (2 * self.num_layers) ** 0.5 |
| 320 | |
| 321 | def forward( |
| 322 | self, |
| 323 | x: torch.Tensor, |
| 324 | freqs_cis: torch.Tensor, |
| 325 | ): |
| 326 | """ |
| 327 | Perform a forward pass through the TransformerBlock. |
| 328 | |
| 329 | Args: |
| 330 | x (torch.Tensor): Input tensor. |
| 331 | freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies. |
| 332 |