1. 项目概述:从零构建Transformer语言模型
斯坦福CS336课程的第一个作业要求我们完整实现一个基于Transformer架构的语言模型(Language Model)。这个作业的核心目标是让学生深入理解现代大语言模型(如GPT系列)的基础构建模块,包括Tokenizer(分词器)、模型架构和优化器等关键组件。
作为一门"从零开始"的课程作业,它与其他NLP课程的最大区别在于:我们需要自己动手实现所有底层组件,而不是简单地调用HuggingFace等现成库。这种教学方式类似于操作系统课程中让学生自己编写操作系统内核,能够让我们真正掌握语言模型的工作原理。
2. 核心组件解析与实现思路
2.1 Tokenizer设计与实现
Tokenizer是将原始文本转换为模型可处理数字序列的关键组件。在作业中,我们需要实现一个基于Byte Pair Encoding (BPE)算法的分词器:
python复制class BPETokenizer:
def __init__(self):
self.vocab = {}
self.merges = []
def train(self, text, vocab_size):
# 初始化词汇表为所有字节
tokens = list(text.encode('utf-8'))
# BPE训练过程
while len(self.vocab) < vocab_size:
# 统计并找到最高频的字节对
pairs = self._get_stats(tokens)
if not pairs:
break
best_pair = max(pairs, key=pairs.get)
# 合并字节对并更新词汇表
self.merges.append(best_pair)
tokens = self._merge(tokens, best_pair)
# 构建最终的词汇表
self.vocab = {i: bytes([i]) for i in range(256)}
# ... (处理合并后的token)
注意:BPE算法需要处理Unicode字符的特殊情况,特别是处理多字节字符时需要考虑字节顺序。
2.2 Transformer架构实现
作业要求实现标准的Transformer解码器架构(类似GPT)。关键组件包括:
python复制def scaled_dot_product_attention(Q, K, V, mask=None):
d_k = Q.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn = torch.softmax(scores, dim=-1)
return torch.matmul(attn, V)
- 多头注意力:
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.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)
def forward(self, x, mask=None):
batch_size = x.size(0)
# 线性变换并分头
Q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
K = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
V = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
# 计算注意力
attn_output = scaled_dot_product_attention(Q, K, V, mask)
# 合并多头输出
attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
return self.W_o(attn_output)
- 前馈网络(FFN):
python复制class FeedForward(nn.Module):
def __init__(self, d_model, d_ff):
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(0.1)
def forward(self, x):
return self.linear2(self.dropout(torch.relu(self.linear1(x))))
2.3 位置编码与层归一化
Transformer需要显式处理序列位置信息。我们使用正弦位置编码:
python复制class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=512):
super().__init__()
position = torch.arange(max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
pe = torch.zeros(max_len, d_model)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe)
def forward(self, x):
return x + self.pe[:x.size(1)]
层归一化(LayerNorm)实现:
python复制class LayerNorm(nn.Module):
def __init__(self, features, eps=1e-6):
super().__init__()
self.gamma = nn.Parameter(torch.ones(features))
self.beta = nn.Parameter(torch.zeros(features))
self.eps = eps
def forward(self, x):
mean = x.mean(-1, keepdim=True)
std = x.std(-1, keepdim=True)
return self.gamma * (x - mean) / (std + self.eps) + self.beta
3. 模型训练与优化
3.1 模型架构整合
完整的Transformer语言模型架构:
python复制class TransformerLM(nn.Module):
def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff):
super().__init__()
self.token_embed = nn.Embedding(vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model)
self.layers = nn.ModuleList([
nn.ModuleDict({
'attn': MultiHeadAttention(d_model, num_heads),
'ffn': FeedForward(d_model, d_ff),
'norm1': LayerNorm(d_model),
'norm2': LayerNorm(d_model)
}) for _ in range(num_layers)
])
self.norm = LayerNorm(d_model)
self.linear = nn.Linear(d_model, vocab_size)
def forward(self, x, mask=None):
x = self.token_embed(x)
x = self.pos_encoding(x)
for layer in self.layers:
# 残差连接和层归一化
attn_output = layer['attn'](layer['norm1'](x), mask)
x = x + attn_output
ffn_output = layer['ffn'](layer['norm2'](x))
x = x + ffn_output
x = self.norm(x)
return self.linear(x)
3.2 优化器实现
作业要求实现Adam优化器。关键步骤包括:
python复制class Adam:
def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8):
self.params = list(params)
self.lr = lr
self.betas = betas
self.eps = eps
self.m = [torch.zeros_like(p) for p in self.params]
self.v = [torch.zeros_like(p) for p in self.params]
self.t = 0
def step(self):
self.t += 1
for i, p in enumerate(self.params):
if p.grad is None:
continue
grad = p.grad.data
# 更新一阶和二阶矩估计
self.m[i] = self.betas[0] * self.m[i] + (1 - self.betas[0]) * grad
self.v[i] = self.betas[1] * self.v[i] + (1 - self.betas[1]) * grad.pow(2)
# 偏差修正
m_hat = self.m[i] / (1 - self.betas[0]**self.t)
v_hat = self.v[i] / (1 - self.betas[1]**self.t)
# 参数更新
p.data -= self.lr * m_hat / (v_hat.sqrt() + self.eps)
3.3 训练流程
完整的训练循环实现:
python复制def train(model, dataloader, optimizer, epochs, device):
model.train()
criterion = nn.CrossEntropyLoss()
for epoch in range(epochs):
total_loss = 0
for batch in dataloader:
inputs, targets = batch
inputs, targets = inputs.to(device), targets.to(device)
# 创建因果掩码
seq_len = inputs.size(1)
mask = torch.tril(torch.ones(seq_len, seq_len)).to(device)
optimizer.zero_grad()
outputs = model(inputs, mask)
loss = criterion(outputs.view(-1, outputs.size(-1)), targets.view(-1))
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f'Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}')
4. 关键问题与解决方案
4.1 内存优化技巧
在实现过程中,内存管理是一个重要挑战。以下是几个关键优化点:
- 梯度检查点(Gradient Checkpointing):
python复制from torch.utils.checkpoint import checkpoint
class TransformerBlock(nn.Module):
def forward(self, x):
# 使用梯度检查点减少内存
return checkpoint(self._forward, x)
def _forward(self, x):
# 实际的前向计算
attn_output = self.attn(self.norm1(x))
x = x + attn_output
ffn_output = self.ffn(self.norm2(x))
return x + ffn_output
- 混合精度训练:
python复制scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
outputs = model(inputs, mask)
loss = criterion(outputs.view(-1, outputs.size(-1)), targets.view(-1))
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
4.2 常见问题排查
- 梯度消失/爆炸:
- 解决方案:使用适当的初始化(如Xavier初始化)、梯度裁剪
python复制torch.nn.init.xavier_uniform_(self.W_q.weight)
torch.nn.init.xavier_uniform_(self.W_k.weight)
torch.nn.init.xavier_uniform_(self.W_v.weight)
# 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
- 训练不稳定:
- 可能原因:学习率过大、未正确使用层归一化
- 解决方案:使用学习率预热、调整层归一化位置
python复制def get_lr(step, d_model, warmup_steps):
return d_model**-0.5 * min(step**-0.5, step * warmup_steps**-1.5)
- 注意力计算效率低:
- 解决方案:实现优化的注意力计算(如FlashAttention)
python复制def flash_attention(Q, K, V, mask=None):
# 分块计算注意力
# ... (实现分块逻辑)
5. 扩展与进阶方向
完成基础实现后,可以考虑以下扩展:
- 模型压缩:
- 知识蒸馏
- 量化感知训练
python复制model = torch.quantization.quantize_dynamic(
model, {nn.Linear}, dtype=torch.qint8
)
- 高效注意力变体:
- 实现稀疏注意力
- 线性注意力
python复制class LinearAttention(nn.Module):
def forward(self, Q, K, V):
# 使用核函数近似softmax
Q = torch.relu(Q)
K = torch.relu(K)
KV = torch.einsum('bhnd,bhne->bhde', K, V)
Z = 1 / (torch.einsum('bhnd,bhd->bhn', Q, K.sum(dim=2)) + self.eps)
return torch.einsum('bhnd,bhde,bhn->bhne', Q, KV, Z)
- 多GPU训练:
python复制model = nn.DataParallel(model, device_ids=[0, 1])
在实现过程中,我发现最关键的insight是:Transformer的核心在于自注意力机制如何动态地建立token之间的关系,而位置编码则为模型提供了必要的序列顺序信息。通过手动实现这些组件,我对现代大语言模型的工作原理有了更深入的理解。
