Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection" This Block has a slightly different structure compared to a regular prenorm Transformer block. The standard block is: LN -> MHA/MLP -> Add. [Ref: https://arxiv.org/abs/2002.0
(
self,
dim,
mixer_cls,
has_text=False,
norm_cls=nn.LayerNorm,
fused_add_norm=False,
residual_in_fp32=False,
drop_path=0.0,
skip=False,
)
| 339 | |
| 340 | class Block(nn.Module): |
| 341 | def __init__( |
| 342 | self, |
| 343 | dim, |
| 344 | mixer_cls, |
| 345 | has_text=False, |
| 346 | norm_cls=nn.LayerNorm, |
| 347 | fused_add_norm=False, |
| 348 | residual_in_fp32=False, |
| 349 | drop_path=0.0, |
| 350 | skip=False, |
| 351 | ): |
| 352 | """ |
| 353 | Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection" |
| 354 | |
| 355 | This Block has a slightly different structure compared to a regular |
| 356 | prenorm Transformer block. |
| 357 | The standard block is: LN -> MHA/MLP -> Add. |
| 358 | [Ref: https://arxiv.org/abs/2002.04745] |
| 359 | Here we have: Add -> LN -> Mixer, returning both |
| 360 | the hidden_states (output of the mixer) and the residual. |
| 361 | This is purely for performance reasons, as we can fuse add and LayerNorm. |
| 362 | The residual needs to be provided (except for the very first block). |
| 363 | """ |
| 364 | super().__init__() |
| 365 | self.residual_in_fp32 = residual_in_fp32 |
| 366 | self.fused_add_norm = fused_add_norm |
| 367 | self.has_text = has_text |
| 368 | self.mixer = mixer_cls(dim) |
| 369 | self.norm = norm_cls(dim) |
| 370 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 371 | if self.fused_add_norm: |
| 372 | assert RMSNorm is not None, "RMSNorm import fails" |
| 373 | assert isinstance( |
| 374 | self.norm, (nn.LayerNorm, RMSNorm) |
| 375 | ), "Only LayerNorm and RMSNorm are supported for fused_add_norm" |
| 376 | self.skip_linear = nn.Linear(2 * dim, dim) if skip else None |
| 377 | |
| 378 | adaln_num = 3 * 2 if self.has_text else 3 |
| 379 | self.adaLN_modulation = nn.Sequential( |
| 380 | nn.SiLU(), nn.Linear(dim, adaln_num * dim, bias=True) |
| 381 | ) |
| 382 | if self.has_text: |
| 383 | self.msa = CrossAttention( |
| 384 | query_dim=dim, context_dim=dim, heads=8, dim_head=64, dropout=0.0 |
| 385 | ) |
| 386 | self.norm_msa = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 387 | |
| 388 | def forward( |
| 389 | self, |