•2025年09月
ai-systems

从零实现 Transformer LLM：PyTorch 自定义分词、多头注意力与生成式训练循环

基于 PyTorch 从零构建 Transformer LLM，涵盖自定义分词、多头注意力及生成训练循环，提供工程参数与最佳实践。

在人工智能领域，大型语言模型（LLM）的兴起彻底改变了自然语言处理范式。从 GPT 系列到 Llama 等模型，它们的核心是 Transformer 架构，而 PyTorch 作为灵活的深度学习框架，是从零实现这些模型的理想工具。本文将一步步指导如何使用 PyTorch 构建一个 Transformer 基础的 LLM，重点聚焦自定义分词、多头注意力机制以及生成式文本建模的训练循环。我们不依赖外部 LLM 库，而是通过纯 PyTorch 代码实现一个小型但功能完整的 GPT-like 模型。这种从头构建的方法不仅能加深对模型内部机制的理解，还能为后续的预训练和微调提供坚实基础。

自定义分词：构建 BPE Tokenizer

分词是 LLM 处理文本的第一步。传统词表方法在处理罕见词和多语言时效率低下，因此 Byte Pair Encoding (BPE) 成为主流选择，如 GPT 模型所采用。BPE 通过迭代合并高频子词对来构建词汇表，既能处理 OOV（Out-of-Vocabulary）问题，又保持词汇表大小在合理范围内（通常 50k-100k）。

在 PyTorch 中实现自定义 BPE tokenizer 的关键是先收集语料库的高频 n-gram，然后训练合并规则。假设我们使用一个小型英文语料库（如 Project Gutenberg 的文本），首先读取并预处理数据：

import re
from collections import defaultdict, Counter

def preprocess_text(text):
    text = re.sub(r'\s+', ' ', text.lower().strip())
    return text.split()

# 加载语料
corpus = []  # 假设从文件加载
words = [word for text in corpus for word in preprocess_text(text)]

接下来，初始化词汇表：将所有词拆分为字符级（包括空格作为特殊 token），统计初始对频率。

# 初始化 pairs
pairs = defaultdict(int)
for word in words:
    symbols = list(word) + ['</w>']  # 词尾标记
    for i in range(len(symbols)-1):
        pairs[(symbols[i], symbols[i+1])] += 1

# 训练 BPE：迭代合并 top 10000 次
num_merges = 10000
vocab = set()
for i in range(num_merges):
    best_pair = max(pairs, key=pairs.get)
    # 合并规则应用到所有词
    new_pairs = defaultdict(int)
    for pair, freq in pairs.items():
        if pair[0] == best_pair[0] and pair[1] == best_pair[1]:
            continue
        # 逻辑合并...
    pairs = new_pairs
    vocab.add(best_pair[0] + best_pair[1])

训练完成后，vocab 包含约 50k 个 subword。编码函数将文本转换为 token ID：

def encode(text, vocab, unk_token='<unk>'):
    words = preprocess_text(text)
    tokens = []
    for word in words:
        symbols = list(word) + ['</w>']
        while len(symbols) > 1:
            pairs = [(symbols[i], symbols[i+1]) for i in range(len(symbols)-1)]
            best = max(pairs, key=lambda p: vocab.get(''.join(p), 0))
            if ''.join(best) in vocab:
                symbols = symbols[:best[0].find(best[0])] + [''.join(best)] + symbols[best[0].find(best[0])+2:]
            else:
                break
        tokens.extend([sym for sym in symbols if sym != '</w>'])
    # 转换为 ID，使用 dict 映射
    return [vocab_id.get(t, unk_id) for t in tokens]

这个实现的关键参数包括：merge 次数（控制 vocab 大小，建议从 8000 开始，根据语料调整至 perplexity 最低）；unk_token 处理率（目标 <1%）；最大序列长度（训练时 1024 tokens）。在实际落地中，使用 Hugging Face 的 tiktoken 作为参考，但自定义版能避免依赖，确保纯 PyTorch 环境。证据显示，这种 BPE 在 GPT 预训练中将 token 效率提升 20%以上，减少了 embedding 层参数量。

数据加载器使用 PyTorch 的 DataLoader，结合自定义 collate_fn 处理变长序列：

from torch.utils.data import Dataset, DataLoader

class TextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=1024):
        self.encodings = [tokenizer.encode(t, max_length) for t in texts]

    def __len__(self):
        return len(self.encodings)

    def __getitem__(self, idx):
        return torch.tensor(self.encodings[idx])

dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

多头注意力机制：Transformer 的核心

Transformer 的威力源于自注意力（Self-Attention），而多头注意力（Multi-Head Attention, MHA）通过并行多个注意力头捕捉不同子空间的依赖关系。假设模型维度 d_model=512，头数 h=8，则每个头 d_k = d_model / h = 64。

PyTorch 实现 MHA 的步骤：

  1. 线性投影：输入 X (batch, seq_len, d_model) 通过 Q、K、V 的线性层。
import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        
        self.scale = math.sqrt(self.d_k)
    
    def forward(self, X, mask=None):
        batch_size, seq_len, _ = X.shape
        
        # 投影
        Q = self.W_q(X).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(X).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(X).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # 注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        if mask is not None:
            scores.masked_fill_(mask == 0, -1e9)
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # 加权求和
        context = torch.matmul(attn_weights, V)
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        
        return self.W_o(context)
  1. 掩码与位置编码：对于因果注意力（decoder-only 如 GPT），使用上三角掩码防止未来信息泄露。位置编码可使用 sin/cos 函数：
def positional_encoding(seq_len, d_model):
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(0, seq_len).unsqueeze(1).float()
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe.unsqueeze(0)

关键参数：num_heads=8（平衡计算与表达力）；dropout=0.1（防止过拟合）；d_model=512（小型模型起点，可扩展至 768）。证据来自 Transformer 原论文，多头设计将 BLEU 分数提升 2-3 点。在实现中，注意 KV 缓存以加速生成（后续章节讨论），并监控注意力权重分布，确保无梯度爆炸（使用 gradient clipping，clip_norm=1.0）。

GPT 模型架构：堆叠 Transformer 块

GPT-like 模型是 decoder-only Transformer，由多个层堆叠而成。每层包括 MHA + 前馈网络（FFN） + 层归一化（LayerNorm）。整体架构：

class GPTBlock(nn.Module):
    def __init__(self, d_model, num_heads, ff_dim=2048, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, ff_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ff_dim, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, X, mask=None):
        attn_out = self.attn(self.norm1(X), mask)
        X = X + self.dropout(attn_out)  # 残差连接
        ff_out = self.ff(self.norm2(X))
        X = X + self.dropout(ff_out)
        return X

class GPTModel(nn.Module):
    def __init__(self, vocab_size, d_model=512, num_layers=6, num_heads=8, max_seq=1024):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_enc = nn.Parameter(positional_encoding(max_seq, d_model))
        self.blocks = nn.ModuleList([GPTBlock(d_model, num_heads) for _ in range(num_layers)])
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        self.d_model = d_model
    
    def forward(self, idx, targets=None):
        batch, seq = idx.shape
        tok_emb = self.embedding(idx) * math.sqrt(self.d_model)
        pos_emb = self.pos_enc[:, :seq, :]
        X = tok_emb + pos_emb
        
        # 因果掩码
        mask = torch.tril(torch.ones(seq, seq)).unsqueeze(0).unsqueeze(0)
        for block in self.blocks:
            X = block(X, mask)
        
        X = self.ln_f(X)
        logits = self.head(X)
        
        if targets is None:
            return logits
        else:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            return logits, loss

参数选择：num_layers=6（小型，参数约 10M）；ff_dim=4*d_model（标准比例）；vocab_size=50257（GPT-2 风格）。这种架构在预训练中以自回归方式预测下一个 token，证据显示层数增加可线性提升 perplexity 降低。

训练循环：预训练与生成式建模

预训练目标是最大化语料似然，使用 AdamW 优化器。完整训练循环：

from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

model = GPTModel(vocab_size)
optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
scheduler = CosineAnnealingLR(optimizer, T_max=1000)  # 总步数

for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        idx = batch[:, :-1]  # 输入
        targets = batch[:, 1:]  # 目标
        _, loss = model(idx, targets)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
    
    print(f"Epoch {epoch}, Loss: {loss.item()}")

关键参数：lr=3e-4（warmup 至 10% 步数后 cosine decay）；batch_size=32（根据 GPU 内存，目标 throughput 10k tokens/s）；warmup_steps=100（线性预热避免初始不稳）；weight_decay=0.1（L2 正则）。监控指标：train_loss（目标 <3.0 for small corpus）；perplexity = exp(loss)（<10 表示好拟合）。生成式文本使用自回归采样：

def generate(model, prompt, max_new=50, temperature=1.0):
    model.eval()
    tokens = torch.tensor(encode(prompt))
    with torch.no_grad():
        for _ in range(max_new):
            logits = model(tokens)[-1, :] / temperature
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, 1)
            tokens = torch.cat([tokens, next_token], dim=1)
    return decode(tokens)  # 反编码

温度=0.8（平衡多样性与连贯）；top_k=50 / top_p=0.9（核采样避免低质输出）。在落地中，使用 KV 缓存加速生成：预计算过去 K/V，避免重复计算，节省 90% 推理时间。

可落地参数与最佳实践

  • 硬件：单 GPU (RTX 3060, 12GB) 足以训练 6-layer 模型；分布式使用 torch.distributed。
  • 超参调优：网格搜索 lr [1e-4, 5e-4]，layers [4-8]；早停于 val_loss 稳定。
  • 风险控制：梯度裁剪防爆炸；混合精度 (torch.amp) 减内存 50%；回滚至 checkpoint 若 loss 上升。
  • 扩展：从此基础微调指令（如 Ch7），或 LoRA 高效 finetune。

通过这些步骤，你能构建一个功能性 LLM，理解从 token 到生成的完整管道。如 Raschka 的仓库所述，这种从零实现是理解 ChatGPT-like 模型的绝佳方式。未来，可扩展至更大语料，实现生产级部署。

（字数约 1250）