1. 项目概述:斯坦福CS336课程作业1解析
这门由斯坦福大学开设的CS336课程"从零开始构建语言模型"堪称大模型时代的"操作系统实践课"。2025年春季学期的第一项作业要求实现Transformer语言模型的核心架构,这正是现代ChatGPT等大语言模型的基石。作为课程的开篇之作,Assignment 1聚焦三大核心组件:分词器(Tokenizer)、模型架构(Model Architecture)和优化器(Optimizer),最终目标是训练出一个最小可用的语言模型。
与常规深度学习课程不同,这门课强调从底层实现所有组件。正如课程描述所言:"代码量至少比其他课程高一个数量级"。作业1看似基础,实则暗藏玄机——你需要在不依赖现成框架的情况下,亲手搭建Transformer的每个齿轮。这种"造轮子"式的训练,正是理解现代语言模型工作原理的最佳途径。
2. 核心组件实现详解
2.1 分词器(Tokenizer)实现
在语言模型中,分词器是将原始文本转化为模型可理解数字序列的第一道关卡。作业要求实现的Byte Pair Encoding(BPE)算法是当前最主流的方案,其核心在于通过迭代合并最高频的字符对来构建词汇表。
实现时需要注意几个关键点:
- 预处理规范化:统一转换为小写、处理unicode、标准化标点
- 合并操作效率:使用优先队列存储字符对频率,O(1)时间获取最高频对
- 特殊token处理:需要保留[UNK]、[PAD]、[BOS]、[EOS]等控制符号
python复制class BPETokenizer:
def __init__(self, corpus, vocab_size):
self.vocab = self.build_vocab(corpus, vocab_size)
def build_vocab(self, corpus, target_size):
# 初始词汇为所有字节字符
vocab = set(corpus)
while len(vocab) < target_size:
# 找出最高频的字符对
pairs = self.get_stats(corpus)
if not pairs:
break
best = max(pairs, key=pairs.get)
# 执行合并操作
corpus = self.merge(corpus, best)
vocab.add(best)
return vocab
2.2 Transformer架构实现
作业要求实现的标准Transformer架构包含以下核心模块:
2.2.1 多头注意力机制
这是Transformer最具创新性的部分。实现时需要特别注意:
- 使用矩阵运算一次性计算所有头的注意力
- 正确实现缩放点积注意力中的mask机制
- 合理处理不同序列长度的padding mask
python复制class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.wq = nn.Linear(d_model, d_model)
self.wk = nn.Linear(d_model, d_model)
self.wv = nn.Linear(d_model, d_model)
self.wo = nn.Linear(d_model, d_model)
def forward(self, q, k, v, mask):
batch_size = q.size(0)
# 线性变换并分头
q = self.wq(q).view(batch_size, -1, self.num_heads, self.head_dim)
k = self.wk(k).view(batch_size, -1, self.num_heads, self.head_dim)
v = self.wv(v).view(batch_size, -1, self.num_heads, self.head_dim)
# 缩放点积注意力
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attention = torch.softmax(scores, dim=-1)
output = torch.matmul(attention, v)
# 合并多头输出
output = output.transpose(1, 2).contiguous()
output = output.view(batch_size, -1, self.d_model)
return self.wo(output)
2.2.2 位置前馈网络(FFN)
虽然结构简单,但实现时要注意:
- 中间层的维度通常是d_model的4倍
- 使用GeLU激活而非ReLU
- 添加适当的dropout防止过拟合
2.2.3 层归一化与残差连接
这两个组件对训练深度Transformer至关重要:
- 层归一化应放在残差连接之前(Pre-LN)
- 残差连接的缩放因子需要谨慎设置
2.3 优化器实现
作业要求实现AdamW优化器,这是训练Transformer的标准选择。与普通Adam的主要区别在于:
- 正确处理权重衰减(weight decay)
- 实现梯度裁剪(gradient clipping)
- 包含偏置校正(bias correction)
python复制class AdamW:
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
self.params = list(params)
self.lr = lr
self.betas = betas
self.eps = eps
self.weight_decay = weight_decay
self.state = {}
def step(self):
for p in self.params:
if p.grad is None:
continue
grad = p.grad.data
if grad.is_sparse:
raise RuntimeError('AdamW does not support sparse gradients')
# 状态初始化
if p not in self.state:
self.state[p] = {
'step': 0,
'exp_avg': torch.zeros_like(p.data),
'exp_avg_sq': torch.zeros_like(p.data)
}
state = self.state[p]
state['step'] += 1
# 更新一阶和二阶动量
state['exp_avg'].mul_(self.betas[0]).add_(grad, alpha=1-self.betas[0])
state['exp_avg_sq'].mul_(self.betas[1]).add_(grad.pow(2), alpha=1-self.betas[1])
# 偏置校正
bias_correction1 = 1 - self.betas[0]**state['step']
bias_correction2 = 1 - self.betas[1]**state['step']
denom = (state['exp_avg_sq'].sqrt() / math.sqrt(bias_correction2)).add_(self.eps)
step_size = self.lr / bias_correction1
# 应用权重衰减
p.data.mul_(1 - self.lr * self.weight_decay)
# 参数更新
p.data.addcdiv_(state['exp_avg'], denom, value=-step_size)
3. 模型训练实战技巧
3.1 数据准备与批处理
训练语言模型需要特别注意数据批处理:
- 动态填充:同一批次内的序列应填充到相同长度
- 注意力掩码:需要为填充位置创建正确的mask
- 批大小选择:根据GPU内存选择合适大小,通常256-1024
python复制def collate_fn(batch):
# 找出批次中最长序列
max_len = max(len(item) for item in batch)
# 初始化输入和mask
inputs = torch.zeros(len(batch), max_len, dtype=torch.long)
masks = torch.zeros(len(batch), max_len, dtype=torch.bool)
for i, item in enumerate(batch):
inputs[i, :len(item)] = torch.tensor(item)
masks[i, :len(item)] = 1
return inputs, masks
3.2 训练循环实现
完整的训练循环包含以下关键步骤:
- 梯度清零
- 前向传播
- 损失计算(交叉熵)
- 反向传播
- 梯度裁剪
- 参数更新
重要提示:在调试阶段,先在CPU上验证代码正确性,再切换到GPU训练。可以设置
torch.autograd.set_detect_anomaly(True)帮助调试。
3.3 超参数设置参考
基于课程建议的初始配置:
- 模型维度(d_model): 512
- 注意力头数: 8
- FFN维度: 2048
- 层数: 6
- 学习率: 3e-4
- 批大小: 256
- Dropout: 0.1
4. 常见问题与调试技巧
4.1 梯度爆炸/消失
症状:损失值变为NaN或剧烈波动
解决方案:
- 检查层归一化的实现
- 添加梯度裁剪(通常1.0-5.0)
- 减小学习率
4.2 模型不收敛
可能原因:
- 分词器实现错误
- 注意力mask不正确
- 学习率设置不当
调试步骤:
- 在极小数据集(如10个样本)上过拟合
- 可视化注意力权重
- 检查token嵌入的梯度
4.3 内存不足
优化策略:
- 使用梯度检查点
- 减少批大小
- 使用混合精度训练
python复制scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
5. 扩展思考与进阶方向
完成基础实现后,可以考虑以下优化:
- 内存效率:实现Flash Attention减少内存占用
- 训练加速:引入混合精度训练
- 模型压缩:尝试知识蒸馏技术
- 推理优化:实现KV缓存加速生成
这门作业虽然只要求实现基础版本,但在实际工业应用中,Transformer的优化空间还很大。例如,可以尝试:
- 不同的注意力变体(如稀疏注意力)
- 替代的归一化方案(RMSNorm)
- 更高效的优化器(如LAMB)
通过这个作业,我深刻体会到现代语言模型虽然接口简单(输入文本,输出文本),但其内部架构的每个组件都经过精心设计。亲手实现一遍后,再看PyTorch的现成实现,会有种"原来如此"的顿悟感。
