File size: 2,474 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import os
import torch
import torch.nn as nn
def _use_reference_rmsnorm() -> bool:
return os.getenv("DIFFULEX_REFERENCE_RMSNORM", "0") == "1"
class RMSNorm(nn.Module):
def __init__(
self,
hidden_size: int,
eps: float = 1e-6,
) -> None:
super().__init__()
self.hidden_size = hidden_size
self.eps = eps
self.weight = nn.Parameter(torch.ones(hidden_size))
@torch.compile
def rms_forward(
self,
x: torch.Tensor,
) -> torch.Tensor:
orig_dtype = x.dtype
x = x.to(torch.float32)
var = x.pow(2).mean(dim=-1, keepdim=True)
x.mul_(torch.rsqrt(var + self.eps))
x = x.to(orig_dtype).mul_(self.weight)
return x
def rms_forward_reference(
self,
x: torch.Tensor,
) -> torch.Tensor:
orig_dtype = x.dtype
x_fp32 = x.to(torch.float32)
var = x_fp32.pow(2).mean(dim=-1, keepdim=True)
x_fp32 = x_fp32 * torch.rsqrt(var + self.eps)
return x_fp32.to(orig_dtype) * self.weight
@torch.compile
def add_rms_forward(
self,
x: torch.Tensor,
residual: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
orig_dtype = x.dtype
x = x.to(torch.float32).add_(residual.to(torch.float32))
residual = x.to(orig_dtype)
var = x.pow(2).mean(dim=-1, keepdim=True)
x.mul_(torch.rsqrt(var + self.eps))
x = x.to(orig_dtype).mul_(self.weight)
return x, residual
def add_rms_forward_reference(
self,
x: torch.Tensor,
residual: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
orig_dtype = x.dtype
x_fp32 = x.to(torch.float32) + residual.to(torch.float32)
residual_out = x_fp32.to(orig_dtype)
var = x_fp32.pow(2).mean(dim=-1, keepdim=True)
x_fp32 = x_fp32 * torch.rsqrt(var + self.eps)
return x_fp32.to(orig_dtype) * self.weight, residual_out
def forward(
self,
x: torch.Tensor,
residual: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
if _use_reference_rmsnorm():
if residual is None:
return self.rms_forward_reference(x)
return self.add_rms_forward_reference(x, residual)
if residual is None:
return self.rms_forward(x)
else:
return self.add_rms_forward(x, residual)
|