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 nn

class 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_embeds
        x = self.drop_emb(x)
        # 通过 Transformer 块
        x = self.trf_blocks(x)
        # 最终归一化
        x = self.final_norm(x)
        # 输出 logits
        logits = 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（高斯误差线性单元）激活函数：

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 = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_resid(x)
        x = x + shortcut
        
        # 前馈网络的残差连接
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_resid(x)
        x = x + shortcut
        
        return 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 = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_resid(x)
        x = x + shortcut
        
        # 前馈网络的残差连接
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_resid(x)
        x = x + shortcut
        
        return 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.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds
        x = 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 的 logits
        logits = 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)