LLM演进史-重建GPT2 -(1)

本文基于 Andrej Karpathy 的 4 小时复现 GPT-2，个人看完后觉得是非常好的视频，这是 LLM 演进史的终结篇，本文基于此进行文字版本的补充。前面的内容请参考 https://blog.nagi.fun 主页的内容，该博主已经写得非常详尽了。

本系列准备一共分为三个部分，分别是主体部分实现，加速实现以及分布式训练。

实现 GPT-2 nn.Module

Config 配置

Copy
@dataclass
class GPTConfig():
block_size: int=1024 # 序列长度限制（上下文窗口长度）
vocab_size: int=50527 # 词表大小
n_layer: int=12 # Transformer层数
n_head: int=12 # 注意力头数量
n_embd: int=768 # 嵌入维度（每个token的向量长度）

@dataclass装饰器定义了一个名为GPTConfig的配置类

(如果不懂装饰器可以自行 CSDN 或者知乎)

为什么要用 dataclass：

•普通类需要手动编写__init__方法，装饰后非常简单
•支持显式声明，且直接pirnt(GPTConfig(n_head=16))可以直接打印出参数

BackBone

Copy

class GPT(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.transformer = nn.ModuleDict(dict(
# word token embedding
wte = nn.Embedding(config.vocab_size, config.n_embd),
# word position embedding
wpe = nn.Embedding(config.block_size, config.n_embd),
# 主体block
h = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
# word token embedding,
ln_f = nn.LayerNorm(config.n_embd),
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.attn = CasualSelfAttention(config)
self.mlp = mlp(config)

def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x

self.transformer：transformer 架构本体

nn.ModuleDict：nn.Module 里面的字典，nn.ModuleDict(ln_f = nn.LayerNorm(config.n_embd) ,)可以理解成{ln_f: nn.LayerNorm(config.n_embd)}

wte ：word 到 embedding 的线性层 [词表的长度，词嵌入维度]，把词变为特征向量

wpe ：word position 到 embedding 的线性层 [序列长度，词嵌入维度]，把位置信息变为特征向量

h ：transformer 的 encoder 主体，每一个 Block 由一个 attention 和一个 mlp 组成

ln_f ：LayerNorm，对于 Pre-Norm 后得到的大方差做一个归一化，后续有解释

lm_head ：最后一层输出，把词的特征向量转为具体的词

Block ：Transformer 由多个相同的 Block 组成

Tips⚠️：此处会发现 GPT2 的 LN 层在 attention 和 mlp 之前，而与上图原文中的 LN 层和残差 (先残差再归一化) 连接在一起不同。

Karpathy的解释是：原本的模型是先残差连接再进行LN归一化处理，这表明连上来的残差也被归一化了，但这是不好的。一个纯净的残差是更好的，因为在反向传播的时候，当梯度回流时，加法是把其梯度均匀分给它的两个分支，这意味着梯度通过残差这条路径直接流向输入，从优化的角度来看，干净的残差是可取的。
说实话我没看懂他的解释，因此去网上搜索了了一下相关的内容，然后发现GPT2的做法叫pre-norm，而Attention is all you need的做法叫post-norm。

苏神对于这两种差异的解释非常到位，残差连接是 x+F(x)x+F(x)x+F(x), 如果 x xx 的方差为 σ 12\sigma^2_1σ12​ ，F(x)F(x)F(x) 的方差为 σ 22\sigma^2_2σ22​，那么残差连接后的方差就是 σ 12+σ 22\sigma^2_1+\sigma^2_2σ12​+σ22​，也就是说残差会放大方差，我们要想办法缩小这个方差，朴素的方法是添加归一化，也就是 x t+1=Norm(xt+F(x))x_{t+1}=Norm(x_t+F(x))xt+1​=Norm(xt​+F(x))，然而，这种做法虽然稳定了前向传播的方差，但事实上已经严重削弱了残差的恒等分支，所以反而失去了残差 “易于训练” 的优点，通常要 warmup 并设置足够小的学习率才能使它收敛。而本身 transformer 的结构有两个特点： warm-up 阶段超参数敏感和优化过程收敛速度慢。(这块笔者也不知道为什么)，相当于在 post-norm 的情况下就更难收敛了，并且训练成本也会有一定程度的上升。

然后来解释一下削弱了残差的恒等分支 (其实也就是 Karpathy 说的干净的残差)，假设初始状态下 x，F(x)x，F(x)x，F(x) 方差都为 1，x+F(x)x+F(x)x+F(x) 方差为 2，而 Normalization 操作负责将方差重新降为 1，这就说明初始阶段 Post Norm 相当

Copy x t+1=xt+F(xt)2 x_{t+1}=\frac{x_t+F(x_t)}{\sqrt{2}}xt+1​=2 ​xt​+F(xt​)​

递归下去

Copy xl=x l−1 2+F l−1(x l−1)2=x l−2 2+F l−2(x l−2)2+F l−1(x l−1)2 x_l=\frac{x_{l-1}}{\sqrt{2}}+\frac{F_{l-1}(x_{l-1})}{\sqrt{2}}=\frac{x_{l-2}}{{2}}+\frac{F_{l-2}(x_{l-2})}{{2}}+\frac{F_{l-1}(x_{l-1})}{\sqrt{2}}xl​=2 ​xl−1​​+2 ​Fl−1​(xl−1​)​=2xl−2​​+2Fl−2​(xl−2​)​+2 ​Fl−1​(xl−1​)​

Copy xl=x 0 2 l/2+F 0(x 0)2 l/2+F 1(x 1)2(l−1)/2+F 2(x 2)2(l−2)/2+...+F l−1(x l−1)2 1/2 x_l=\frac{x_{0}}{{2^{l/2}}}+\frac{F_{0}(x_{0})}{{2^{l/2}}}+\frac{F_{1}(x_{1})}{{2}^{(l-1)/2}}+\frac{F_{2}(x_{2})}{{2}^{(l-2)/2}}+...+\frac{F_{l-1}(x_{l-1})}{{2}^{1/2}}xl​=2l/2x0​​+2l/2F0​(x0​)​+2(l−1)/2F1​(x1​)​+2(l−2)/2F2​(x2​)​+...+21/2Fl−1​(xl−1​)​

本来残差的意思是给前面输入的层搞一条“绿色通道”，让梯度可以更直接地回传，但是在Post Norm中，这条“绿色通道”被严重削弱了，越靠近前面的通道反而权重越小，使得在多次残差连接后，前面的残差无法感知到末尾的梯度变化，残差“名存实亡”，因此还是不容易训练。
论文请见《ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE》
而修正后的 Pre-Norm 其形式为

Copy x t+1=xt+Ft(Norm(xt))x_{t+1}=x_t+F_t(Norm(x_t))xt+1​=xt​+Ft​(Norm(xt​))

展开迭代后：

Copy x t+1=xt+Ft(Norm(xt))=x t−1+F t−1(Norm(x t−1))+Ft(Norm(xt))x_{t+1}=x_t+F_t(Norm(x_t)) = x_{t-1}+F_{t-1}(Norm(x_{t-1}))+F_t(Norm(x_t))xt+1​=xt​+Ft​(Norm(xt​))=xt−1​+Ft−1​(Norm(xt−1​))+Ft​(Norm(xt​))

Copy xt=x 0+F 0(Norm(x 0))+F 1(Norm(x 1))+...+F l−1(Norm(x l−1))x_{t}=x_{0}+F_{0}(Norm(x_{0}))+F_1(Norm(x_1))+...+F_{l-1}(Norm(x_{l-1}))xt​=x0​+F0​(Norm(x0​))+F1​(Norm(x1​))+...+Fl−1​(Norm(xl−1​))

每一条残差通道都是平权的，残差的作用会比 Post Norm 更加明显，所以它也更好优化。当然，这样最后的 xl x_lxl​方差将会很大，所以在接预测层之前 xl x_lxl​ 也还要加个 Normalization，这正是ln_f

Karpathy 讲到 Attention 是 tokens 进行通信的地方，是一个池化函数，是一个加权和函数，是一个 ruduce operation 。
MLP 发生在每个单独的 token 上，在 tokens 之间没有信息被收集或者交换，是一个 map operation 。

MLP

Copy
class mlp(nn.Module):
def __init__(self, config):
super().__init__()
self.c_fc = nn.Linear(config.n_embd, config.n_embd*4)
self.Gelu = nn.GELU(approximate='tanh')
self.c_proj = nn.Linear(config.n_embd*4, config.n_embd)

def forward(self, x):
x = self.c_fc(x)
x = self.Gelu(x)
x = self.c_proj(x)
return x

很简单的 MLP 线性映射，从 [n_embd, 4 * n_embd] 到 [4 * n_embd, n_embd] ，中间再加一个非线性层 GELU 激活函数。GELU 的函数图像和 Relu 很像，但是在尾部处可导，也就是解决了 Relu 在 x 小于 0 时候导数恒为 0 的问题，而这种平滑会产生更好的效果

Karpathy 这里讲到为什么使用 tanh 的近似值，谈到这是一个历史遗留问题，在 tensorflow 时期用精确的 GELU 特别慢，所以开发了利用 tanh 近似 GELU 的函数。

Attention

Copy
class CasualSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
self.c_attn = nn.Linear(config.n_embd, config.n_embd*3)
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
self.n_embd = config.n_embd
self.n_head = config.n_head
# 后续在权重模块会讲到它的作用
self.c_proj.NANOGPT_SCALE_INIT = 1
self.register_buffer("bias", torch.tril((torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size)))

def forward(self, x):
B, T, C = x.size()
qkv = self.c_attn(x)
# 得到qkv
q, k, v = qkv.split(self.n_embd, dim=2)
# query, key, value 全部被分成[B, n_head, T, n_embd//n_head]
query = q.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
key = k.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
value = v.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
# QK^T/d
att = query @ key.transpose(-1, -2) * (1.0/math.sqrt(key.size(-1)))
mask_att = att.masked_fill(self.bias[:,:,:T,:T]==0, float('-inf'))
wei = F.softmax(mask_att, dim=-1)
out = wei @ value
out = out.transpose(1,2).contiguous().view(B, T, C)
out = self.c_proj(out)
return out

self.c_attn： Wq,Wk,Wv W_q,W_k,W_vWq​,Wk​,Wv​ 的组合，将输入 x xx 变为输入的 Q,K,V Q,K,VQ,K,V

self.c_proj：计算完 Q KT dk V\frac{QK^T}{\sqrt{d_k}}Vdk​ ​QKT​V 后的一层线性层

self.n_embd：每个 token 的特征向量空间

self.n_head：多头注意力机制的头的个数

self.bias：这里的 bias 是掩码的意思，也就是上三角矩阵，防止前面的 token 学习到后面 token 的方法。具体原理如下：对于输入 x xx :

Copy x=[x 11x 12x 13x 21x 22x 23x 31x 32x 33],bias=[100110111],mas k att=[x 11−inf−infx 21x 22−infx 31x 32x 33]x=\begin{bmatrix} x_{11} & x_{12} & x_{13} \\ x_{21} & x_{22} & x_{23} \\ x_{31} & x_{32} & x_{33} \end{bmatrix},bias=\begin{bmatrix} 1 & 0 & 0 \\ 1 & 1 & 0 \\ 1 & 1 & 1 \end{bmatrix} ,mask_{att}=\begin{bmatrix} x_{11} & -inf & -inf \\ x_{21} & x_{22} & -inf \\ x_{31} & x_{32} & x_{33} \end{bmatrix}x= ​x11​x21​x31​​x12​x22​x32​​x13​x23​x33​​ ​,bias= ​111​011​001​ ​,maskatt​= ​x11​x21​x31​​−infx22​x32​​−inf−infx33​​ ​

-inf 在后续的softmax 过程中会变成接近于 0 的一个值，从而对分类不起影响

contiguous()：transpose 并不改变物理排序，只会改变形式排序，而使用该函数才可以修正物理排序，举个例子说明，对于数组 [[[1,2][7,8]],[[3,4][5,6]]] shape=[1,2,2][[[1,2][7,8]],[[3,4][5,6]]]\space \space \space \space \space shape=[1,2,2][[[1,2][7,8]],[[3,4][5,6]]] shape=[1,2,2]， transpose(1,2)后等于[[[1,2][2,3]],[[3,5][4,6]]] shape=[1,2,2][[[1,2][2,3]],[[3,5][4,6]]]\space \space \space \space \space shape=[1,2,2][[[1,2][2,3]],[[3,5][4,6]]] shape=[1,2,2]，但是这两个数组的物理存储都为 [1,2,7,8,3,4,5,6][1,2,7,8,3,4,5,6][1,2,7,8,3,4,5,6] ，因此在对 transpose 后的数组进行 view 操作时会报错。

DownLoad from HuugingFace

Copy
# 在 class GPT 下
@classmethod
def from_pretrained(cls, model_type):
"""Loads pretrained GPT-2 model weights from huggingface"""
# 四种类型的model
assert model_type in {'gpt2','gpt2-medium','gpt2-large','gpt2-xl'}
# 打印出你加载的是哪一种
print("Loading weights from pretrained gpt:%s"%model_type)
# 每种GPT所对应的超参数不一样大
config_args ={
'gpt2' : dict(n_layer=12,n_head=12,n_embd=768), # 124M params
'gpt2-medium' : dict(n_layer=24,n_head=16,n_embd=1024), #350M params
'gpt2-large' : dict(n_layer=36,n_head=20,n_embd=1280), #774M params
'gpt2-xl' : dict(n_layer=48,n_head=25,n_embd=1600), #1558M params
}[model_type]
# 词表大小永远是50527
config_args['vocab_size'] = 50257
# 单个block的大小始终是1024
config_args['block_size'] = 1024
# 给模型导入超参数
config = GPTConfig(**config_args)
model = GPT(config)
# sd是模型参数名字典
sd = model.state_dict()
sd_keys = sd.keys()
sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask

# 从HF上下载权重，sd_hf是下载下来的模型参数名字典
model_hf = GPT2LMHeadModel.from_pretrained(model_type, cache_dir="/home/shong_Tan/project/gpt_2/model_weight", local_files_only=True)
sd_hf = model_hf.state_dict()

sd_keys_hf = sd_hf.keys()
# 丢掉HF权重中掩码bias
sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')]
# 丢掉HF权重中bias
sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')]
transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
# 确保sd和hf_sd的参数名一样多
assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
# 确保sd和hf_sd的transformer块的权重名一样
for k in sd_keys_hf:
if any(k.endswith(w) for w in transposed):
assert sd_hf[k].shape[::-1] == sd[k].shape
with torch.no_grad():
sd[k].copy_(sd_hf[k].t())
else:
assert sd_hf[k].shape == sd[k].shape, f"mismatched keys: {sd_hf[k].shape} != {sd[k].shape}"
with torch.no_grad():
sd[k].copy_(sd_hf[k])
return model

看代码注释就好了

Tips⚠️：从 HF 下载下来的lm_head.weight 和 transformer.wte.weight 的形状大小是一样的，都是 [50527,768][50527, 768][50527,768]，它们一个是输入嵌入，一个是输出 logits，这两个应该是一致的，这才能反应出当 token 被嵌入特征向量时的语义在经过交互后，当输出的时候还是这个特征向量，它又可以被转回到原来的 token。同时， 50527∗768≈40 M 50527*768 \approx 40M50527∗768≈40M，可以节省大量的显存。

Forward

Copy
# 在 class GPT 下
def forward(self, idx, target):
# 进入时维度为[batch, token长度]
B, T = idx.size()
# token长度不能超过上下文
assert T <= self.config.block_size, f"超出输入上下文长度限制 {T-self.config.block_size} token"
# pos [0,1,2,..,T-1]，并且记住要放到device上去
pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
# 位置嵌入
pos = self.transformer.wpe(pos) #(T, n_embd)
# token嵌入
tok = self.transformer.wte(idx) #(B, T, n_embd)
# 相加是在 (T, n_embd)这个维度上的数值相加
x = tok + pos
# 经过transformer块
for block in self.transformer.h:
x = block(x)
# 最后一层归一化
x = self.transformer.ln_f(x)
# 线性层输出
logits = self.lm_head(x) #(B, T, vocab_size)
loss = None
# 有target的话也就是有标签就进行训练，计算损失函数，反之只用推理就可以
if target is not None:
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1))
return logits, loss

Copy
# 小试牛刀
num_return_sequences = 5
max_length = 30
model = GPT.from_pretrained('gpt2')
# eval会在评估的时候禁用dropout层，对于batchnorm也会有不一样的反应，以及冻结参数
model.eval()
# 模型移到gpu
model.to('cuda')

# 以下是分词，调openai的tiktoken库就行了，如果想要知道原理的，建议去看文章开头的博主blog
import tiktoken
enc = tiktoken.get_encoding('gpt2')
tokens = enc.encode("Hello, I'm a language model, ")
tokens = torch.tensor(tokens, dtype=torch.long) # [8, ]
tokens = tokens.unsqueeze(0).repeat(num_return_sequences, 1) # [5,8]
x = tokens.to('cuda')

while x.size(1) < max_length:
with torch.no_grad():
# 输入模型得到结果
logits, loss = model(x) # x: [B, T] logits:[B,T,C]
# 取出最后一个token的预测
logits = logits[:, -1, :] # [B, 1, C]
# 对最后一个维度C取softmax
probs = F.softmax(logits, dim=-1) # [B, 1, C]
# 在最后一个维度C选出topk大的概率和对应的index
topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
# 从topk的概率中随机选取一个概率
ix = torch.multinomial(topk_probs, 1)
# 找到被选中的概率对应的index
xcol = torch.gather(topk_indices, -1, ix)
# 将得到的输出token加在x上又作为输入 [B, T+1]
x = torch.cat((x, xcol), dim=1)

分词形式：将 "Hello, I'm a language model," 变为 [15496, 11, 314, 1101, 257, 3303, 2746, 11, 220]

参考下面网站可以自己尝试

https://tiktokenizer.vercel.app/

初始化

Dataset

Copy
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
# 这是mac的M芯片系列
elif hasattr(torch.backends, "eps") and torch.backends.mps.is_available():
device = 'mps'
print("Using device: ", device)
import tiktoken
enc = tiktoken.get_encoding('gpt2')
with open('input.txt', 'r') as f:
data = f.read()
text = data[:1000]
tokens = enc.encode(text)
B, T = 4, 32
buf = torch.tensor(tokens[:B*T+1])
buf.to(device)
# 本质上是通过前n个词推测n+1个词
x = buf[:-1].view(B, T)
y = buf[1:].view(B,T)

model.GPT(GPTConfig())
model.to(device)
logits, loss = model(x)
print(loss.item())

这里loss的值大概是 11 左右，因为−log(150257)≈11-log(\frac{1}{50257})\approx 11−log(502571​)≈11

训练单个 batch 代码

Copy
# 使用AdamW优化器，Adam与SGD的区别自己去了解吧
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(50):
# 优化器梯度清零
optimizer.zero_grad()
# 得到logits和损失
logits, loss = model(x, y)
# 反向传播求导
loss.backward()
# 利用导数更新原有参数
optimizer.step()

Adam 优化器能比 SGD 更快地收敛

Dataloader 函数

Copy
class DataLoaderLite():

def __init__(self, B, T):
self.B = B
self.T = T
# 读入整个input.txt
with open('input.txt', 'r') as f:
data = f.read()
enc = tiktoken.get_encoding('gpt2')
tokens = enc.encode(data)
self.tokens = torch.tensor(tokens, dtype=torch.long)
print(f"load {len(self.tokens)} tokens")
print(f"1 epoch = {len(self.tokens)//(B*T)} batched")
# 定义在当前batch的位置
self.current_position = 0

def next_batch(self):
B, T = self.B, self.T
buf = self.tokens[self.current_position: self.current_position+B*T+1]
x = buf[:-1].view(B, T)
y = buf[1:].view(B, T)
# 每个batch有B*T个元素对
self.current_position += B*T
# 如果batch把tokens用完了，又回到toknes[0]
if self.current_position+B*T+1 > len(self.tokens):
self.current_position = 0
return x, y

修正训练代码

Copy
train_loader = DataLoaderLite(4, 32 )
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(50):
optimizer.zero_grad()
x, y = train_loader.next_batch()
x, y = x.to(device), y.to(device)
logits, loss = model(x, y)
loss.backward()
optimizer.step()
print(f"step: {i}, loss: {loss.item()}")

权重

Copy
def _init_weights(self, module):
if isinstance(module, nn.Linear):
std = 0.02
if hasattr(module, ' '):
std = std * (2*self.config.n_layer**-0.5)
torch.nn.init.normal_(module.weight, mean=0, std=std)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)

elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0, std=0.02)

std = std * (2*self.config.n_layer)**-0.5 ：此处的方差来源是考虑了残差流的贡献，每次残差连接都说明输入 input 有了一次等额贡献，需要一个因子来处理， 1 2∗n layer\frac{1}{\sqrt{2*n_{layer}}}2∗nlayer​ ​1​，这里是控制了 Pre-Norm 的残差连接导致的方差过大，这里的 2 是因为每层里面 Attention 和 MLP 都用到了一次残差。

std ：std 的值的来源也是有根据的，按照 GPT2 里文档的说法，最好在 1 n embd\frac{1}{\sqrt{n_{embd}}}nembd​ ​1​ 左右