ereniko commited on
Commit
724d0e3
·
verified ·
1 Parent(s): a726540

Delete rmsnorm.py

Browse files
Files changed (1) hide show
  1. rmsnorm.py +0 -23
rmsnorm.py DELETED
@@ -1,23 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
-
4
-
5
- class RMSNorm(nn.Module):
6
- """RMSNorm (Section 4.7): cheaper alternative to LayerNorm.
7
-
8
- Rescales by root-mean-square of the activations instead of full
9
- mean/variance normalization. No bias, single learnable scale per dim.
10
- """
11
-
12
- def __init__(self, dim: int, eps: float = 1e-5):
13
- super().__init__()
14
- self.eps = eps
15
- self.weight = nn.Parameter(torch.ones(dim))
16
-
17
- def forward(self, x: torch.Tensor) -> torch.Tensor:
18
- # compute in float32 for stability regardless of input dtype (bf16 etc.)
19
- dtype = x.dtype
20
- x = x.float()
21
- rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
22
- out = x * rms
23
- return (out.to(dtype)) * self.weight