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

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

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

实现 GPT-2 nn.Module#

Config 配置#

@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#


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

transformer

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。

pre-norm

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

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

x_{t+1}=\frac{x_t+F(x_t)}{\sqrt{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}}

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}}

本来残差的意思是给前面输入的层搞一条“绿色通道”，让梯度可以更直接地回传，但是在Post Norm中，这条“绿色通道”被严重削弱了，越靠近前面的通道反而权重越小，使得在多次残差连接后，前面的残差无法感知到末尾的梯度变化，残差“名存实亡”，因此还是不容易训练。

论文请见《ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE》

而修正后的 Pre-Norm 其形式为

x_{t+1}=x_t+F_t(Norm(x_t))

展开迭代后：

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))

x_{t}=x_{0}+F_{0}(Norm(x_{0}))+F_1(Norm(x_1))+...+F_{l-1}(Norm(x_{l-1}))

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

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

MLP#

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 的函数。

GELU

Attention#

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： $W_q,W_k,W_v$ 的组合，将输入 $x$ 变为输入的 $Q,K,V$

self.c_proj：计算完 $\frac{QK^T}{\sqrt{d_k}}V$ 后的一层线性层

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

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

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

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}

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

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

DownLoad from HuugingFace#

    # 在 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]$ ，它们一个是输入嵌入，一个是输出 logits，这两个应该是一致的，这才能反应出当 token 被嵌入特征向量时的语义在经过交互后，当输出的时候还是这个特征向量，它又可以被转回到原来的 token。同时， $50527*768 \approx 40M$ ，可以节省大量的显存。

Forward#

    # 在 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

# 小试牛刀
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#

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(\frac{1}{50257})\approx 11$

训练单个 batch 代码

# 使用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 函数

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

修正训练代码

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()}")

权重#

    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 有了一次等额贡献，需要一个因子来处理， $\frac{1}{\sqrt{2*n_{layer}}}$ ，这里是控制了 Pre-Norm 的残差连接导致的方差过大，这里的 2 是因为每层里面 Attention 和 MLP 都用到了一次残差。

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