| 353 | |
| 354 | |
| 355 | class TransformerBlock(nn.Module): |
| 356 | def __init__(self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int, |
| 357 | multiple_of: int, ffn_dim_multiplier: float, norm_eps: float, |
| 358 | qk_norm: bool) -> None: |
| 359 | """ |
| 360 | Initialize a TransformerBlock. |
| 361 | |
| 362 | Args: |
| 363 | layer_id (int): Identifier for the layer. |
| 364 | dim (int): Embedding dimension of the input features. |
| 365 | n_heads (int): Number of attention heads. |
| 366 | n_kv_heads (Optional[int]): Number of attention heads in key and |
| 367 | value features (if using GQA), or set to None for the same as |
| 368 | query. |
| 369 | multiple_of (int): Value to ensure hidden dimension is a multiple |
| 370 | of this value in the FeedForward block. |
| 371 | ffn_dim_multiplier (float, optional): Custom multiplier for hidden |
| 372 | dimension in the FeedForward block. Defaults to None. |
| 373 | norm_eps (float): A small value added to the norm layer |
| 374 | denominators to avoid division-by-zero. |
| 375 | |
| 376 | Attributes: |
| 377 | n_heads (int): Number of attention heads. |
| 378 | dim (int): Dimension size of the model. |
| 379 | head_dim (int): Dimension size of each attention head. |
| 380 | attention (Attention): Attention module. |
| 381 | feed_forward (FeedForward): FeedForward module. |
| 382 | layer_id (int): Identifier for the layer. |
| 383 | attention_norm (RMSNorm): Layer normalization for attention output. |
| 384 | ffn_norm (RMSNorm): Layer normalization for feedforward output. |
| 385 | adaLN_modulation (nn.Sequential): A small network to generate |
| 386 | feature modulation factors. |
| 387 | |
| 388 | """ |
| 389 | super().__init__() |
| 390 | self.dim = dim |
| 391 | self.head_dim = dim // n_heads |
| 392 | self.attention = Attention(dim, n_heads, n_kv_heads, qk_norm) |
| 393 | self.feed_forward = FeedForward( |
| 394 | dim=dim, hidden_dim=4 * dim, multiple_of=multiple_of, |
| 395 | ffn_dim_multiplier=ffn_dim_multiplier, |
| 396 | ) |
| 397 | self.layer_id = layer_id |
| 398 | self.attention_norm = RMSNorm(dim, eps=norm_eps) |
| 399 | self.ffn_norm = RMSNorm(dim, eps=norm_eps) |
| 400 | |
| 401 | self.adaLN_modulation = nn.Sequential( |
| 402 | nn.SiLU(), |
| 403 | ColumnParallelLinear( |
| 404 | min(dim, 1024), 6 * dim, bias=True, gather_output=True, |
| 405 | init_method=nn.init.zeros_, |
| 406 | ), |
| 407 | ) |
| 408 | |
| 409 | def forward( |
| 410 | self, |
| 411 | x: torch.Tensor, |
| 412 | freqs_cis: torch.Tensor, |