Views
No views yet
1from transformers import AutoTokenizer, GPT2Model
2
3import torch.nn as nn
4import torch
5
6class RMSLayerNorm(nn.Module):
7 def __init__(self, normalized_shape, eps=1e-8, affine=True):
8 super(RMSLayerNorm, self).__init__()
9 self.normalized_shape = normalized_shape
10 self.eps = eps
11 self.affine = affine
12
13 if self.affine:
14 self.weight = nn.Parameter(torch.ones(()))
15 else:
16 self.register_parameter('weight', None)
17 self.register_parameter('bias', None)
18
19 def forward(self, x):
20 rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
21 x_normalized = x / rms
22 if self.affine:
23 x_normalized = x_normalized * self.weight
24 return x_normalized
25
26
27def replace(model):
28 for name, child in model.named_children():
29 if isinstance(child, nn.modules.normalization.LayerNorm):
30 setattr(model, name, RMSLayerNorm(child.normalized_shape, eps=child.eps, affine=True))
31 else:
32 replace(child)
33 return model
34
35
36class GPTR2Model(GPT2Model):
37 def __init__(self, config):
38 super().__init__(config)
39 replace(self)
40
41model = GPTR2Model.from_pretrained("George-Ogden/gptr2-nano-without-momentum-with-weight-decay")
42tokenizer = AutoTokenizer.from_pretrained("gpt2")