在人工智能领域，大型语言模型（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 = defaultdict(int)
for word in words:
symbols = list(word) + ['</w>']
for i in range(len(symbols)-1):
pairs[(symbols[i], symbols[i+1])] += 1
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>'])
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 的威力源于自注意力（Self-Attention），而多头注意力（Multi-Head Attention, MHA）通过并行多个注意力头捕捉不同子空间的依赖关系。假设模型维度 d_model=512，头数 h=8，则每个头 d_k = d_model / h = 64。
PyTorch 实现 MHA 的步骤：
- 线性投影：输入 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)
- 掩码与位置编码：对于因果注意力（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-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）