LLM基础5_从零开始实现 GPT 模型
基于GitHub项目:https://github.com/datawhalechina/llms-from-scratch-cn
设计 LLM 的架构
GPT 模型基于 Transformer 的 decoder-only 架构,其主要特点包括:
-
顺序生成文本
-
参数数量庞大(而非代码量复杂)
-
大量重复的模块化组件
以 GPT-2 small 模型(124M 参数)为例,其配置如下:
GPT_CONFIG_124M = {"vocab_size": 50257, # BPE 分词器词表大小"ctx_len": 1024, # 最大上下文长度"emb_dim": 768, # 嵌入维度"n_heads": 12, # 注意力头数量"n_layers": 12, # Transformer 块层数"drop_rate": 0.1, # Dropout 比例"qkv_bias": False # QKV 计算是否使用偏置
}
GPT 模型基本结构
cfg是配置实例
import torch.nn as nnclass GPTModel(nn.Module):def __init__(self, cfg):super().__init__()# Token 嵌入层self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])# 位置嵌入层self.pos_emb = nn.Embedding(cfg["ctx_len"], cfg["emb_dim"])# Dropout 层self.drop_emb = nn.Dropout(cfg["drop_rate"])# 堆叠n_layers相同的Transformer 块self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])# 最终层归一化self.final_norm = LayerNorm(cfg["emb_dim"])# 输出层self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)def forward(self, in_idx):batch_size, seq_len = in_idx.shape# Token 嵌入tok_embeds = self.tok_emb(in_idx)# 位置嵌入pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))# 组合嵌入x = tok_embeds + pos_embedsx = self.drop_emb(x)# 通过 Transformer 块x = self.trf_blocks(x)# 最终归一化x = self.final_norm(x)# 输出 logitslogits = self.out_head(x)return logits
层归一化 (Layer Normalization)
层归一化将激活值规范化为均值为 0、方差为 1 的分布,加速模型收敛:
class LayerNorm(nn.Module):def __init__(self, emb_dim):super().__init__()self.eps = 1e-5 # 防止除零错误的标准设定值self.scale = nn.Parameter(torch.ones(emb_dim)) #可学习缩放参数,初始化为全1向量self.shift = nn.Parameter(torch.zeros(emb_dim)) #可学习平移参数,初始化为全0向量def forward(self, x):mean = x.mean(dim=-1, keepdim=True) #计算均值 μ,沿最后一维,保持维度var = x.var(dim=-1, keepdim=True, unbiased=False) #计算方差 σ²,同均值维度,有偏估计(分母n)norm_x = (x - mean) / torch.sqrt(var + self.eps) #标准化计算,分母添加ε防溢出return self.scale * norm_x + self.shift #仿射变换,恢复模型表达能力
GELU 激活函数与前馈网络
GPT 使用 GELU(高斯误差线性单元)激活函数:
场景 | ReLU 的行为 | GELU 的行为 |
---|---|---|
处理微弱负信号 | 直接丢弃(可能丢失细节) | 部分保留(如:保留 30% 的信号强度) |
遇到强烈正信号 | 完全放行 | 几乎完全放行(保留 95% 以上) |
训练稳定性 | 容易在临界点卡顿 | 平滑过渡,减少训练震荡 |
应对复杂模式 | 需要堆叠更多层数 | 单层就能捕捉更细腻的变化 |
class GELU(nn.Module):def __init__(self):super().__init__()def forward(self, x):return 0.5 * x * (1 + torch.tanh(torch.sqrt(torch.tensor(2.0 / torch.pi)) * (x + 0.044715 * torch.pow(x, 3))))
前馈神经网络实现:
class FeedForward(nn.Module):def __init__(self, cfg):super().__init__()self.layers = nn.Sequential(nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),GELU(),nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),nn.Dropout(cfg["drop_rate"]))def forward(self, x):return self.layers(x)
Shortcut 连接
Shortcut 连接(残差连接)解决深度网络中的梯度消失问题:
class TransformerBlock(nn.Module):def __init__(self, cfg):super().__init__()self.att = MultiHeadAttention(d_in=cfg["emb_dim"],d_out=cfg["emb_dim"],block_size=cfg["ctx_len"],num_heads=cfg["n_heads"], dropout=cfg["drop_rate"],qkv_bias=cfg["qkv_bias"])self.ff = FeedForward(cfg)self.norm1 = LayerNorm(cfg["emb_dim"])self.norm2 = LayerNorm(cfg["emb_dim"])self.drop_resid = nn.Dropout(cfg["drop_rate"])def forward(self, x):# 注意力块的残差连接shortcut = xx = self.norm1(x)x = self.att(x)x = self.drop_resid(x)x = x + shortcut# 前馈网络的残差连接shortcut = xx = self.norm2(x)x = self.ff(x)x = self.drop_resid(x)x = x + shortcutreturn x
Transformer 块整合
将多头注意力与前馈网络整合为 Transformer 块:
class TransformerBlock(nn.Module):def __init__(self, cfg):super().__init__()self.att = MultiHeadAttention(d_in=cfg["emb_dim"],d_out=cfg["emb_dim"],block_size=cfg["ctx_len"],num_heads=cfg["n_heads"], dropout=cfg["drop_rate"],qkv_bias=cfg["qkv_bias"])self.ff = FeedForward(cfg)self.norm1 = LayerNorm(cfg["emb_dim"])self.norm2 = LayerNorm(cfg["emb_dim"])self.drop_resid = nn.Dropout(cfg["drop_rate"])def forward(self, x):# 注意力块的残差连接shortcut = xx = self.norm1(x)x = self.att(x)x = self.drop_resid(x)x = x + shortcut# 前馈网络的残差连接shortcut = xx = self.norm2(x)x = self.ff(x)x = self.drop_resid(x)x = x + shortcutreturn x
完整 GPT 模型实现
class GPTModel(nn.Module):def __init__(self, cfg):super().__init__()self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])self.pos_emb = nn.Embedding(cfg["ctx_len"], cfg["emb_dim"])self.drop_emb = nn.Dropout(cfg["drop_rate"])self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])self.final_norm = LayerNorm(cfg["emb_dim"])self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)def forward(self, in_idx):batch_size, seq_len = in_idx.shapetok_embeds = self.tok_emb(in_idx)pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))x = tok_embeds + pos_embedsx = self.drop_emb(x)x = self.trf_blocks(x)x = self.final_norm(x)logits = self.out_head(x)return logits
文本生成
使用贪婪解码生成文本:
def generate_text_simple(model, idx, max_new_tokens, context_size):for _ in range(max_new_tokens):# 截断超过上下文长度的部分idx_cond = idx[:, -context_size:]with torch.no_grad():logits = model(idx_cond)# 获取最后一个 token 的 logitslogits = logits[:, -1, :] probas = torch.softmax(logits, dim=-1)idx_next = torch.argmax(probas, dim=-1, keepdim=True)idx = torch.cat((idx, idx_next), dim=1)return idx
使用示例:
# 初始化模型
model = GPTModel(GPT_CONFIG_124M)# 设置评估模式
model.eval()# 生成文本
start_context = "Every effort moves you"
encoded = tokenizer.encode(start_context)
encoded_tensor = torch.tensor(encoded).unsqueeze(0)generated = generate_text_simple(model=model,idx=encoded_tensor,max_new_tokens=10,context_size=GPT_CONFIG_124M["ctx_len"]
)decoded_text = tokenizer.decode(generated.squeeze(0).tolist())
print(decoded_text)