| 570 | # TODO: (Dhruv) This can be replaced with regular RMSNorm in Mochi once `_keep_in_fp32_modules` is supported |
| 571 | # for sharded checkpoints, see: https://github.com/huggingface/diffusers/issues/10013 |
| 572 | class MochiRMSNorm(nn.Module): |
| 573 | def __init__(self, dim, eps: float, elementwise_affine: bool = True): |
| 574 | super().__init__() |
| 575 | |
| 576 | self.eps = eps |
| 577 | |
| 578 | if isinstance(dim, numbers.Integral): |
| 579 | dim = (dim,) |
| 580 | |
| 581 | self.dim = torch.Size(dim) |
| 582 | |
| 583 | if elementwise_affine: |
| 584 | self.weight = nn.Parameter(torch.ones(dim)) |
| 585 | else: |
| 586 | self.weight = None |
| 587 | |
| 588 | def forward(self, hidden_states): |
| 589 | input_dtype = hidden_states.dtype |
| 590 | variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) |
| 591 | hidden_states = hidden_states * torch.rsqrt(variance + self.eps) |
| 592 | |
| 593 | if self.weight is not None: |
| 594 | hidden_states = hidden_states * self.weight |
| 595 | hidden_states = hidden_states.to(input_dtype) |
| 596 | |
| 597 | return hidden_states |
| 598 | |
| 599 | |
| 600 | class GlobalResponseNorm(nn.Module): |
no outgoing calls
no test coverage detected
searching dependent graphs…