r"""Pass the input through the encoder layer. Args: hidden_states: the sequence to the encoder layer (required). residual: hidden_states = Mixer(LN(residual))
(
self,
x: Tensor,
residual: Optional[Tensor] = None,
c=None,
text=None,
inference_params=None,
skip=None,
)
| 386 | self.norm_msa = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 387 | |
| 388 | def forward( |
| 389 | self, |
| 390 | x: Tensor, |
| 391 | residual: Optional[Tensor] = None, |
| 392 | c=None, |
| 393 | text=None, |
| 394 | inference_params=None, |
| 395 | skip=None, |
| 396 | ): |
| 397 | r"""Pass the input through the encoder layer. |
| 398 | |
| 399 | Args: |
| 400 | hidden_states: the sequence to the encoder layer (required). |
| 401 | residual: hidden_states = Mixer(LN(residual)) |
| 402 | """ |
| 403 | if self.skip_linear is not None: |
| 404 | x = self.skip_linear(torch.cat([x, skip], dim=-1)) |
| 405 | |
| 406 | if not self.fused_add_norm: |
| 407 | if residual is None: |
| 408 | residual = x |
| 409 | else: |
| 410 | residual = residual + self.drop_path(x) |
| 411 | |
| 412 | x = self.norm(residual.to(dtype=self.norm.weight.dtype)) |
| 413 | if self.residual_in_fp32: |
| 414 | residual = residual.to(torch.float32) |
| 415 | else: |
| 416 | fused_add_norm_fn = ( |
| 417 | rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn |
| 418 | ) |
| 419 | if residual is None: |
| 420 | x, residual = fused_add_norm_fn( |
| 421 | x, |
| 422 | self.norm.weight, |
| 423 | self.norm.bias, |
| 424 | residual=residual, |
| 425 | prenorm=True, |
| 426 | residual_in_fp32=self.residual_in_fp32, |
| 427 | eps=self.norm.eps, |
| 428 | ) |
| 429 | else: |
| 430 | x, residual = fused_add_norm_fn( |
| 431 | self.drop_path(x), |
| 432 | self.norm.weight, |
| 433 | self.norm.bias, |
| 434 | residual=residual, |
| 435 | prenorm=True, |
| 436 | residual_in_fp32=self.residual_in_fp32, |
| 437 | eps=self.norm.eps, |
| 438 | ) |
| 439 | |
| 440 | if not self.has_text: |
| 441 | shift_mba, scale_mba, gate_mba = self.adaLN_modulation(c).chunk(3, dim=1) |
| 442 | x = x + gate_mba.unsqueeze(1) * self.mixer( |
| 443 | modulate(x, shift_mba, scale_mba), |
| 444 | inference_params=inference_params, |
| 445 | ) |