Spaces:
Sleeping
Sleeping
| from typing import Any | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| class TransposeLast(nn.Module): | |
| def __init__( | |
| self, | |
| deconstruct_idx: int | None = None, | |
| transpose_dim: int = -2, | |
| ) -> None: | |
| super().__init__() | |
| self.deconstruct_idx = deconstruct_idx | |
| self.transpose_dim = transpose_dim | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| if self.deconstruct_idx is not None: | |
| x = x[self.deconstruct_idx] | |
| return x.transpose(self.transpose_dim, -1) | |
| class Fp32LayerNorm(nn.LayerNorm): | |
| def forward(self, input: torch.Tensor) -> torch.Tensor: | |
| output = F.layer_norm( | |
| input.float(), | |
| self.normalized_shape, | |
| self.weight.float() if self.weight is not None else None, | |
| self.bias.float() if self.bias is not None else None, | |
| self.eps, | |
| ) | |
| return output.type_as(input) | |
| class Fp32GroupNorm(nn.GroupNorm): | |
| use_triton: bool = False | |
| def forward( | |
| self, | |
| input: torch.Tensor, | |
| padding_mask: torch.BoolTensor | None = None, | |
| ) -> torch.Tensor: | |
| if self.use_triton: | |
| from src.modeling.kernels.group_norm import masked_group_norm | |
| return masked_group_norm(input, padding_mask, self.weight, self.bias, self.num_groups, self.eps) | |
| if padding_mask is not None: | |
| batch_size, channels, sequence_length = input.shape | |
| grouped = input.float().view( | |
| batch_size, | |
| self.num_groups, | |
| channels // self.num_groups, | |
| sequence_length, | |
| ) | |
| valid = (~padding_mask).float()[:, None, None, :] | |
| count = valid.sum(dim=-1, keepdim=True) * grouped.shape[2] | |
| mean = (grouped * valid).sum(dim=(2, 3), keepdim=True) / count | |
| variance = ((grouped - mean).square() * valid).sum(dim=(2, 3), keepdim=True) / count | |
| output = ((grouped - mean) * torch.rsqrt(variance + self.eps)).view_as(input) | |
| if self.affine: | |
| output = output * self.weight.float()[None, :, None] | |
| output = output + self.bias.float()[None, :, None] | |
| return output.type_as(input) | |
| output = F.group_norm( | |
| input.float(), | |
| self.num_groups, | |
| self.weight.float() if self.weight is not None else None, | |
| self.bias.float() if self.bias is not None else None, | |
| self.eps, | |
| ) | |
| return output.type_as(input) | |
| class GradMultiply(torch.autograd.Function): | |
| def forward(ctx: Any, x: torch.Tensor, scale: float) -> torch.Tensor: | |
| ctx.scale = scale | |
| return x.new(x) | |
| def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None]: | |
| return grad * ctx.scale, None | |