Lecture 06: GPT 系列

核心主题：GPT 架构实现、从零预训练、Flash Attention、分布式训练 (DDP)、HellaSwag 评测

1. GPT 模型家族

1.1 GPT-2 模型规格

2. 文本生成策略

2.1 最佳实践配置

generation_config = {
    "top_k": 50,
    "top_p": 0.95,
    "temperature": 0.7,
    "repetition_penalty": 1.2,
}

3. GPT-2 架构实现 ⭐⭐

3.1 与原始 Transformer 的区别

3.2 核心组件代码

GPTConfig 配置：

from dataclasses import dataclass

@dataclass
class GPTConfig:
    block_size: int = 1024      # 最大序列长度
    vocab_size: int = 50257     # GPT-2 词汇表大小
    n_layer: int = 12           # Transformer 层数
    n_head: int = 12            # 注意力头数
    n_embd: int = 768           # 嵌入维度

CausalSelfAttention (含 Flash Attention)：

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # Q, K, V 投影合并为一个线性层
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        # 输出投影
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.c_proj.NANOGPT_SCALE_INIT = 1  # 残差流缩放标记
        self.n_head = config.n_head
        self.n_embd = config.n_embd

    def forward(self, x):
        B, T, C = x.size()  # batch, seq_len, embedding_dim
        # 计算 Q, K, V
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)
        # 重塑为多头: (B, n_head, T, head_dim)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        # Flash Attention (PyTorch 2.0+)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        # 合并多头
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        # 输出投影
        y = self.c_proj(y)
        return y

Block (Pre-Norm 结构)：

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.gelu = nn.GELU(approximate='tanh')
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
        self.c_proj.NANOGPT_SCALE_INIT = 1

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        return x

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        # Pre-Norm: LN → Sublayer → Residual Add
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

GPT 模型：

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),   # token embedding
            wpe = nn.Embedding(config.block_size, config.n_embd),   # position embedding
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),  # 最终 LayerNorm
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        # Weight Tying: 共享 token embedding 和 output 权重
        self.transformer.wte.weight = self.lm_head.weight
        # 初始化参数
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            # 残差流路径上的投影层缩放初始化
            if hasattr(module, 'NANOGPT_SCALE_INIT'):
                std *= (2 * self.config.n_layer) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.config.block_size
        # 前向传播
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        tok_emb = self.transformer.wte(idx)      # (B, T, n_embd)
        pos_emb = self.transformer.wpe(pos)      # (T, n_embd)
        x = tok_emb + pos_emb
        # Transformer blocks
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)                 # (B, T, vocab_size)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

4. 关键训练技巧

4.1 Weight Tying ⭐

原理：

• $W_e \in \mathbb{R}^{V \times d}$：将 token ID 映射为嵌入向量
• $W_{out} \in \mathbb{R}^{d \times V}$：将隐状态映射回词汇表
• 令 $W_{out} = W_e^T$，则语义相近的 token 在输出空间也相近

参数节省：

占 GPT-2 Small (124M) 总参数的约 31%！

# Weight Tying 实现
self.transformer.wte.weight = self.lm_head.weight
# 反向传播时，两处梯度会自动累加到同一张量

4.2 残差流缩放初始化 (Residual Stream Scaling)

残差路径上的投影层（c_proj）使用缩放标准差初始化：

动机：每个 block 有 2 条残差路径（attention + MLP），共 $2N$ 次累加。缩放确保残差流的方差不会随深度爆炸。

# 标记残差路径上的投影层
self.c_proj.NANOGPT_SCALE_INIT = 1

# 初始化时应用缩放
if hasattr(module, 'NANOGPT_SCALE_INIT'):
    std *= (2 * self.config.n_layer) ** -0.5

4.3 Vocab Size 对齐

GPT-2 原始词汇表大小为 50257（不是 2 的幂次），修改为 50304：

5. Flash Attention ⭐

5.1 标准注意力的问题

标准自注意力需要显式计算并存储完整的注意力矩阵：

内存瓶颈：

• 注意力矩阵 $QK^T$ 大小为 $O(T^2)$
• 序列长度 T=1024 时：$1024 \times 1024 \times 4\text{bytes} \approx 4\text{MB/head}$
• T=4096 时：$\approx 64\text{MB/head}$，多头 + batch 后显存占用巨大

IO 瓶颈：

• $QK^T$ 在 HBM 中分配 → 写入 HBM → 读回计算 softmax → 再写入 → 再读回乘 V
• 多次 HBM 读写是主要性能瓶颈（GPU 计算速度远快于显存带宽）

5.2 Flash Attention 优势

关键技术：

1. Tiling (分块)：将 Q, K, V 分成小块，逐块在 SRAM 中计算
2. Kernel Fusion (算子融合)：将 matmul → scale → mask → softmax → matmul 融合为单一 CUDA kernel
3. Online Softmax：增量计算 softmax，无需完整行数据
4. Recomputation：反向传播时重新计算注意力（trade compute for memory）

# PyTorch 2.0+ 原生支持 Flash Attention
# 自动根据硬件选择最优实现 (FlashAttention-2, Memory-Efficient, Math)
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# 等价于手动实现但速度提升 2-4x，显存节省 5-20x

6. 完整预训练配置 (FineWeb-Edu 10B)

6.1 超参数表

6.2 Cosine LR with Warmup

import math

def get_lr(step, warmup_steps, max_steps, max_lr, min_lr):
    # 1) 线性 warmup 阶段
    if step < warmup_steps:
        return max_lr * (step + 1) / warmup_steps
    # 2) 训练结束后保持最小学习率
    if step > max_steps:
        return min_lr
    # 3) Cosine decay 阶段
    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # 从1衰减到0
    return min_lr + coeff * (max_lr - min_lr)

6.3 优化器配置 (Weight Decay 分离)

def configure_optimizers(model, weight_decay, learning_rate, device):
    # 收集所有需要梯度的参数
    param_dict = {pn: p for pn, p in model.named_parameters() if p.requires_grad}
    # 区分需要 weight decay 的参数 (2D 权重矩阵) 和不需要的 (bias, LayerNorm)
    decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
    nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
    optim_groups = [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ]
    num_decay = sum(p.numel() for p in decay_params)
    num_nodecay = sum(p.numel() for p in nodecay_params)
    print(f"decay params: {num_decay:,}, no-decay params: {num_nodecay:,}")
    # 使用 fused AdamW (更快的 CUDA 实现)
    use_fused = 'cuda' in device
    optimizer = torch.optim.AdamW(
        optim_groups, lr=learning_rate,
        betas=(0.9, 0.95), eps=1e-8,
        fused=use_fused
    )
    return optimizer

6.4 训练优化清单

6.5 DDP 分布式训练 (Distributed Data Parallel)

启动命令：

# 单机 8 GPU
torchrun --standalone --nproc_per_node=8 train_gpt2.py

# 多机 (2 nodes x 8 GPUs)
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
         --master_addr=host0 --master_port=29500 train_gpt2.py

DDP 初始化与梯度同步：

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化进程组
dist.init_process_group(backend='nccl')
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
device = f'cuda:{ddp_local_rank}'
torch.cuda.set_device(device)
master_process = (ddp_rank == 0)  # 主进程负责 logging/checkpoint

# 包装模型
model = GPT(GPTConfig(vocab_size=50304))
model = model.to(device)
model = DDP(model, device_ids=[ddp_local_rank])
raw_model = model.module  # 访问原始模型

# 训练循环中的梯度同步控制
for micro_step in range(grad_accum_steps):
    # 最后一步才同步梯度，前面的步骤不需要 all-reduce
    if micro_step < grad_accum_steps - 1:
        model.require_backward_grad_sync = False
    else:
        model.require_backward_grad_sync = True
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        logits, loss = model(x, y)
    loss = loss / grad_accum_steps  # 归一化
    loss.backward()

7. HellaSwag 评测

7.1 任务描述

HellaSwag (Harder Endings, Longer contexts, and Low-shot Activities for Situations With Adversarial Generations) 是一个常识推理 benchmark：

• 给定一个场景描述（context），从 4 个候选续写中选择最合理的一个
• 对抗性生成的干扰项，对人类很简单（~95%）但对模型很难
• 测试模型的常识推理和语言理解能力

7.2 评测方法代码

def get_most_likely_row(tokens, mask, logits):
    """
    评测逻辑：计算每个候选续写的平均 token 对数概率，
    选择概率最高的作为模型预测。
    """
    # 将 logits 向左移一位对齐 (next token prediction)
    shift_logits = (logits[..., :-1, :]).contiguous()
    shift_tokens = (tokens[..., 1:]).contiguous()
    # 计算每个位置的 log probability
    flat_shift_logits = shift_logits.view(-1, shift_logits.size(-1))
    flat_shift_tokens = shift_tokens.view(-1)
    shift_losses = F.cross_entropy(
        flat_shift_logits, flat_shift_tokens, reduction='none'
    )
    shift_losses = shift_losses.view(tokens.size(0), -1)
    # 只关注续写部分 (mask=1 的位置)
    shift_mask = (mask[..., 1:]).contiguous()
    masked_shift_losses = shift_losses * shift_mask
    # 计算每个候选的平均 loss
    sum_loss = masked_shift_losses.sum(dim=1)
    avg_loss = sum_loss / shift_mask.sum(dim=1)
    # loss 最小 = 概率最大 → 模型的选择
    pred_idx = avg_loss.argmin().item()
    return pred_idx

7.3 结果对比

8. 从验证到全流程

8.1 单 batch 过拟合验证

目的：在正式训练前验证模型实现的正确性。

# 验证流程：用单个 batch 反复训练，loss 应降到接近 0
# 初始 loss 应接近 -ln(1/50257) ≈ 10.82
# Step 0: loss = 10.96 (随机初始化，略高于理论值)
# Step 1: loss = 9.22
# ...
# Step 50: loss = 0.01 (成功过拟合 → 模型实现正确)

8.2 随机模型 vs 预训练模型对比

验证流程总结：

1. 检查初始 loss $\approx \ln(V)$，确认初始化正确
2. 单 batch 过拟合，确认前向/反向/优化器工作正常
3. 对比 HuggingFace 预训练权重，确认架构兼容
4. 启动完整训练，观察 loss 曲线和 HellaSwag 提升

核心要点总结