一、课程概述与目标

本课程是Andrej Karpathy「0 to Hero」系列的核心章节，目标是从零复现GPT-2 124M参数模型——从加载OpenAI预训练权重验证模型正确性，到完全随机初始化参数、基于高质量数据集从零训练，最终实现性能对标甚至超越原版GPT-2。

核心价值在于：

  • 掌握Decoder-only Transformer的完整实现（Pre-LN结构、权重共享等GPT-2关键特性）
  • 理解大模型训练的核心技术栈（混合精度、梯度累积、分布式训练、性能优化）
  • 学会从数据处理、模型初始化到评估全流程的工程实践

二、模型结构核心解析（GPT-2本质）

GPT-2是典型的Decoder-only Transformer，与原始Transformer的核心差异的是「Pre-LN」结构和权重共享，模型整体骨架如下：

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

class GPTConfig:
    vocab_size = 50257  # GPT-2固定词汇表大小
    block_size = 1024   # 最大序列长度
    n_layer = 12        # Transformer层数量
    n_head = 12         # 注意力头数
    n_embd = 768        # 嵌入维度
    dropout = 0.0       # GPT-2默认无dropout
    bias = False        # 输出层无偏置

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # 1. 核心组件：token嵌入 + 位置嵌入
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),  # token嵌入层
            wpe = nn.Embedding(config.block_size, config.n_embd),  # 位置嵌入层
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),  # Transformer层
            ln_f = nn.LayerNorm(config.n_embd),  # 输出前的层归一化（Pre-LN新增）
        ))
        
        # 2. 输出层（语言模型头）：与token嵌入层权重共享
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=config.bias)
        
        # 权重共享核心代码：lm_head与wte共享权重
        self.transformer.wte.weight = self.lm_head.weight
        
        # 3. 初始化权重
        self.apply(self._init_weights)
        # 残差层权重缩放（GPT-2关键初始化策略）
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))

    def _init_weights(self, module):
        # GPT-2初始化规则：线性层用0.02标准差正态分布，偏置置0
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            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, "序列长度超过模型最大限制"
        
        # 生成位置索引：0~T-1
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        
        # 嵌入层：token嵌入 + 位置嵌入
        tok_emb = self.transformer.wte(idx)  # (B, T, n_embd)
        pos_emb = self.transformer.wpe(pos)  # (T, n_embd) → 广播到(B, T, n_embd)
        x = tok_emb + pos_emb
        
        # 经过所有Transformer Block
        for block in self.transformer.h:
            x = block(x)
        
        # 输出层归一化 + 线性投影
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)  # (B, T, vocab_size)
        
        # 计算损失（若传入targets）
        loss = None
        if targets is not None:
            # 展平logits和targets以适配cross_entropy
            B, T, C = logits.size()
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)
        
        return logits, loss

关键结构要点

  1. Pre-LN结构：层归一化（LayerNorm）放在Attention/MLP之前，而非之后。优势是保持残差流纯净，梯度传播更顺畅。
  2. 权重共享：lm_head与wte共享同一套权重矩阵（shape: 50257×768），既减少30%参数量（节省40M参数），又让输入输出的语义空间对齐。
    • 核心逻辑：嵌入层是「token→向量」，输出层是「向量→token相似度打分」，本质都是同一语义空间的映射，无需单独学习两套权重。
  3. 位置嵌入：GPT-2的位置嵌入是可学习参数（而非正弦编码），训练中会自动学习位置相关的特征。

三、核心组件实现（Attention/MLP/Block）

1. 多头自注意力（Multi-Head Attention）

GPT-2的Attention是「自回归掩码注意力」，确保每个token只能看到前面的token，同时通过多头机制捕捉不同维度的依赖关系：

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0, "嵌入维度必须能被头数整除"
        
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)  # QKV投影
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)        # 输出投影
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        
    def forward(self, x):
        B, T, C = x.size()
        
        # 1. QKV投影：(B, T, 3C) → 拆分Q/K/V各为(B, T, C)
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)
        
        # 2. 多头拆分：(B, T, C) → (B, n_head, T, C//n_head)
        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)
        
        # 3. 自回归掩码注意力（使用PyTorch原生Flash Attention加速）
        x = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            dropout_p=0.0,
            is_causal=True  # 自动生成下三角掩码，实现自回归
        )
        
        # 4. 多头合并：(B, n_head, T, C//n_head) → (B, T, C)
        x = x.transpose(1, 2).contiguous().view(B, T, C)
        
        # 5. 输出投影
        x = self.c_proj(x)
        return x

2. 前馈网络（MLP）

GPT-2的MLP使用「GELU近似激活」，由两个线性层和激活函数组成，作用是对每个token进行独立的非线性变换（无跨token信息交互）：

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)  # 升维
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)  # 降维
        self.gelu = nn.GELU(approximate='tanh')  # GPT-2使用近似版GELU
    
    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        return x

3. Transformer Block

每个Block由「多头自注意力 + 残差连接 + MLP + 残差连接」组成，Pre-LN结构的核心体现：

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, bias=config.bias)  # Attention前归一化
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd, bias=config.bias)  # MLP前归一化
        self.mlp = MLP(config)
    
    def forward(self, x):
        # 残差连接：x + Attention(LayerNorm(x))
        x = x + self.attn(self.ln_1(x))
        # 残差连接：x + MLP(LayerNorm(x))
        x = x + self.mlp(self.ln_2(x))
        return x

四、数据处理：从数据集到训练批次

1. 数据集选择与tokenization

  • 调试阶段：使用tiny Shakespeare（100万字符，适合快速验证代码）
  • 正式训练：使用FineWeb Edu（1000亿高质量教育类token，过滤后噪声低）

tokenization使用GPT-2原生tokenizer（tiktoken）：

import tiktoken

# 加载GPT-2 tokenizer
tokenizer = tiktoken.get_encoding("gpt2")

def encode(text):
    return tokenizer.encode(text, allowed_special={"<|endoftext|>"})

def decode(tokens):
    return tokenizer.decode(tokens)

# 示例：编码文本
text = "Hello I'm a language model"
tokens = encode(text)  # [15496, 1004, 16, 13, 318, 2746]

2. 批次构造核心：为什么需要T+1？

语言模型的目标是「预测下一个token」，因此需要从连续token序列中同时构造「输入序列」和「目标序列」：

  • 输入序列（x）：前T个token
  • 目标序列（y）：后T个token
  • 因此需要先截取T+1长度的token序列，再拆分：
def create_batch(tokens, B, T):
    # tokens: 一维token序列（长度≥B*T + 1）
    buf = torch.tensor(tokens[:B*T + 1], dtype=torch.long)  # 取B*T+1个token
    x = buf[:-1].view(B, T)  # 输入：(B, T) → 前B*T个token
    y = buf[1:].view(B, T)   # 目标：(B, T) → 后B*T个token
    return x, y

# 示例：B=4, T=6（批次大小4，序列长度6）
tokens = encode("Hello I'm a language model. I can generate coherent text.")
x, y = create_batch(tokens, 4, 6)
print(x.shape)  # torch.Size([4, 6])
print(y.shape)  # torch.Size([4, 6])

3. 数据加载器实现（支持分布式训练）

class DataLoaderLite:
    def __init__(self, path, B, T, rank=0, world_size=1):
        self.B = B
        self.T = T
        self.rank = rank
        self.world_size = world_size
        
        # 加载并编码文本
        with open(path, 'r', encoding='utf-8') as f:
            text = f.read()
        self.tokens = torch.tensor(encode(text), dtype=torch.long)
        
        # 分布式数据分片：每个进程处理不同的起始位置
        self.current_pos = self.rank * B * T
    
    def next_batch(self):
        # 截取当前批次的token（B*T + 1个）
        buf = self.tokens[self.current_pos : self.current_pos + self.B*self.T + 1]
        x, y = create_batch(buf.numpy(), self.B, self.T)
        
        # 更新位置，循环遍历数据
        self.current_pos += self.B * self.T * self.world_size
        if self.current_pos + self.B*self.T + 1 > len(self.tokens):
            self.current_pos = self.rank * self.B * self.T
        
        return x, y

五、训练核心技术：稳定与性能优化

1. 初始化策略：为什么需要权重缩放？

GPT-2的初始化包含两个关键步骤，与层归一化分工明确：

  • 权重缩放：解决「初始化时方差爆炸」，对残差分支（Attention/MLP的输出投影层）的权重按1/√(2*n_layer)缩放
    • 原理：深度残差网络中，每层残差分支的输出方差会累积，缩放后可确保初始化时方差为1
  • 层归一化：解决「训练中分布偏移」，在每个子层前标准化输入，稳定梯度传播

2. 混合精度训练：BF16与FP32的分工

混合精度训练的核心是「在不损失精度的前提下提升速度/节省显存」，GPT-2中的精度分配：

数据类型 适用场景 原因
BF16（半精度） 前向/反向传播的中间计算（矩阵乘法、Attention、MLP） 计算密集型操作对精度不敏感，BF16可减少显存占用50%，提升GPU利用率
FP32（单精度） 模型参数、损失计算、梯度更新、Softmax/层归一化 对精度敏感，避免数值不稳定（如梯度消失、损失计算偏差）

实现代码（PyTorch autocast）：

# 混合精度训练上下文
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    logits, loss = model(x, targets=y)
    loss = loss / grad_accum_steps  # 梯度累积时归一化损失

# 反向传播（梯度仍以FP32存储）
loss.backward()

3. 性能优化三件套：torch.compile + Flash Attention + vocab_size对齐

（1）torch.compile：一键加速

PyTorch 2.0+的编译器，通过算子融合、消除Python开销实现加速（2~3倍提升）：

# 编译模型（仅需一行代码）
model = torch.compile(model)

核心优化：将多个element-wise操作（如GELU+线性层）融合为单个GPU kernel，减少显存读写开销。

（2）Flash Attention：注意力计算加速

替代传统Attention实现，通过「不 materialize 注意力矩阵」减少显存占用，提升30%+速度：

# 直接使用PyTorch原生Flash Attention（已集成在CausalSelfAttention中）
x = F.scaled_dot_product_attention(q, k, v, is_causal=True)

原理：利用在线Softmax技巧，避免存储T×T的注意力矩阵，大幅减少显存读写。

（3）vocab_size对齐：避免低效边界计算

GPT-2默认vocab_size=50257（奇数，非2的幂），GPU kernel对2的幂数优化更好，因此调整为53040（最近的2的幂倍数）：

# 修改配置：对齐vocab_size为2的幂
config = GPTConfig(vocab_size=53040)

优势：减少GPU kernel的边界处理开销，提升4%~30%速度（PyTorch 2.3+效果更明显）。

4. 优化器配置：AdamW + 权重衰减

GPT-2使用AdamW优化器，权重衰减仅作用于二维参数（矩阵类权重），一维参数（偏置、LayerNorm缩放因子）不衰减：

def configure_optimizers(model, weight_decay=0.1, lr=6e-4, betas=(0.9, 0.95), eps=1e-8):
    # 拆分衰减/不衰减参数
    decay_params = []
    nodecay_params = []
    for n, p in model.named_parameters():
        if p.dim() >= 2:  # 二维参数（线性层权重、嵌入层）：衰减
            decay_params.append(p)
        else:  # 一维参数（偏置、LayerNorm）：不衰减
            nodecay_params.append(p)
    
    optim_groups = [
        {'params': decay_params, 'weight_decay': weight_decay},
        {'params': nodecay_params, 'weight_decay': 0.0}
    ]
    
    # 使用fused AdamW（更快的GPU实现）
    optimizer = torch.optim.AdamW(optim_groups, lr=lr, betas=betas, eps=eps, fused=True)
    return optimizer

5. 梯度累积：模拟大批次训练

当GPU显存不足以容纳大批次时，通过「多次小批次累积梯度」模拟大批次效果：

# 配置
total_batch_size = 524288  # 目标大批次（2^19）
B = 64  # 单GPU微批次
T = 1024  # 序列长度
world_size = 8  # GPU数量
grad_accum_steps = total_batch_size // (B * T * world_size)  # 梯度累积步数

# 训练循环中的梯度累积逻辑
for step in range(max_steps):
    loss_accum = 0.0
    for micro_step in range(grad_accum_steps):
        x, y = train_loader.next_batch()
        x, y = x.to(device), y.to(device)
        
        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
            logits, loss = model(x, targets=y)
            loss = loss / grad_accum_steps  # 归一化损失
            loss_accum += loss.detach()
        
        loss.backward()  # 累积梯度
    
    # 梯度裁剪（最后一道稳定防线）
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    
    optimizer.step()
    optimizer.zero_grad()

6. 全局梯度裁剪：防止梯度爆炸

与初始化、层归一化的分工：

  • 初始化：解决「训练前」的方差问题
  • 层归一化：解决「训练中」的分布偏移问题
  • 梯度裁剪：解决「参数更新前」的梯度爆炸问题

实现代码：

# 裁剪所有参数的全局梯度范数
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

核心逻辑：

  1. 计算所有参数梯度的L2范数：||g||₂ = √(sum(g_i²))
  2. 若||g||₂ > max_norm，则缩放所有梯度：g_i' = g_i × (max_norm / ||g||₂)
  3. 保持梯度方向不变，仅限制步长

7. 学习率调度：Warmup + 余弦衰减

遵循GPT-3的学习率策略，避免训练初期学习率过大导致不稳定：

def get_lr(step, warmup_steps=715, max_steps=19073, max_lr=6e-4, min_lr=6e-5):
    # 1. 线性warmup：从0提升到max_lr
    if step < warmup_steps:
        return max_lr * (step + 1) / warmup_steps
    # 2. 余弦衰减：从max_lr衰减到min_lr
    progress = (step - warmup_steps) / (max_steps - warmup_steps)
    lr = min_lr + 0.5 * (max_lr - min_lr) * (1 + torch.cos(torch.tensor(progress * math.pi)))
    return lr.item()

# 训练循环中更新学习率
for step in range(max_steps):
    lr = get_lr(step)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

8. 分布式训练（DDP）：多GPU协同

利用PyTorch DDP实现多GPU并行训练，核心是「数据分片 + 梯度同步」：

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

def setup_ddp():
    # 初始化DDP
    dist.init_process_group(backend='nccl')
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    torch.cuda.set_device(rank)
    return rank, world_size

def main():
    rank, world_size = setup_ddp()
    is_master = rank == 0  # 主进程负责日志和保存模型
    
    # 数据加载（每个GPU处理不同分片）
    train_loader = DataLoaderLite(path='input.txt', B=64, T=1024, rank=rank, world_size=world_size)
    
    # 模型初始化并移到GPU
    model = GPT(GPTConfig()).to(rank)
    model = DDP(model, device_ids=[rank])  # 包装DDP
    
    # 优化器
    optimizer = configure_optimizers(model.module)  # 注意用model.module访问原始模型
    
    # 训练循环（与单GPU一致，DDP自动处理梯度同步）
    for step in range(max_steps):
        # ... 训练逻辑不变 ...
        
        # 仅主进程打印日志
        if is_master and step % 100 == 0:
            print(f"step {step}, loss: {loss_accum.item()}, lr: {lr}")

if __name__ == "__main__":
    main()

运行命令（8个GPU）：

torchrun --standalone --nproc_per_node=8 train_gpt2.py

六、评估：验证模型性能

1. 验证损失

定期在验证集上计算损失，监控模型泛化能力：

def evaluate(model, val_loader, device):
    model.eval()
    total_loss = 0.0
    n_steps = 20
    with torch.no_grad():
        for _ in range(n_steps):
            x, y = val_loader.next_batch()
            x, y = x.to(device), y.to(device)
            logits, loss = model(x, targets=y)
            total_loss += loss.item()
    model.train()
    return total_loss / n_steps

2. Hellaswag评估：常识推理能力

Hellaswag是常识推理基准，通过「句子补全」任务评估模型的世界知识：

  • 任务形式：给上下文 + 4个候选结尾，选最合理的一个
  • 实现逻辑：将4个候选结尾构造为4个序列，计算每个序列的平均损失，损失最低的为预测结果
def evaluate_hellaswag(model, device, num_examples=1000):
    model.eval()
    correct = 0
    with torch.no_grad():
        for example in load_hellaswag_examples(num_examples):
            # 构造4个候选序列（上下文 + 每个候选结尾）
            tokens, mask, label = render_hellaswag_example(example)
            tokens = tokens.to(device)
            mask = mask.to(device)
            
            # 计算logits
            logits, _ = model(tokens)
            
            # 计算每个候选的平均损失
            losses = []
            for i in range(4):
                # 只计算候选部分的损失
                candidate_tokens = tokens[i:i+1]
                candidate_mask = mask[i:i+1]
                target = candidate_tokens[:, 1:]
                logit = logits[i:i+1, :-1]
                loss = F.cross_entropy(
                    logit.reshape(-1, logit.size(-1)),
                    target.reshape(-1),
                    reduction='none'
                )
                # 按mask加权平均
                avg_loss = (loss * candidate_mask[:, 1:].reshape(-1)).sum() / candidate_mask.sum()
                losses.append(avg_loss)
            
            # 选损失最小的候选
            pred = torch.argmin(torch.tensor(losses))
            if pred == label:
                correct += 1
    
    model.train()
    return correct / num_examples * 100

七、训练结果与总结

1. 训练配置

  • 数据集：FineWeb Edu（100亿token）
  • 硬件：8×A100 GPU（80GB）
  • 训练时长：1.7小时（单epoch）/ 8小时（4 epoch）
  • 性能指标：
    • 吞吐量：150万token/秒
    • Hellaswag准确率：33.24%（超越GPT-2原版的29.55%）
    • 验证损失：低于GPT-2原版

2. 核心收获

  1. 模型层面：掌握GPT-2的Pre-LN结构、权重共享、可学习位置嵌入等关键特性。
  2. 训练层面：理解初始化、层归一化、梯度裁剪的分工，混合精度、梯度累积、分布式训练的工程实现。
  3. 优化层面：学会用torch.compile、Flash Attention、vocab_size对齐等技巧提升性能。

3. 后续优化方向

  • 数据层面：对数据集进行随机打乱，减少训练中的周期性波动。
  • 模型层面：增加dropout正则化，防止过拟合。
  • 评估层面：集成更多基准（如MMLU、GSM8K）全面评估模型能力。
  • 应用层面：通过监督微调（SFT）将预训练模型转化为对话模型。

附录：完整代码目录

gpt2-reproduction/
├── train_gpt2.py        # 主训练脚本（含模型、数据加载、训练循环）
├── hellaswag.py         # Hellaswag评估脚本
├── data_loader.py       # 数据加载器实现
└── config.py            # GPT-2配置参数
# 深度学习 # pytorch # 人工智能 # python
Logo
2048 AI社区

有“AI”的1024 = 2048，欢迎大家加入2048 AI社区

更多推荐

AI原生应用领域内容生成的数据分析与洞察

我们的目的是深入了解在AI原生应用这个神奇的领域里，内容生成背后的数据分析与洞察到底是怎么回事。我们会从基本概念开始，逐步探索算法原理、实际应用等方面，范围涵盖了从理论知识到实际操作的各个环节，就像一次精彩的探险之旅，带大家全面认识这个有趣的领域。接下来的内容会按照一定的顺序展开。首先我们会介绍一些核心概念，让大家对AI原生应用、内容生成、数据分析与洞察有个初步的认识；然后讲解核心算法原理和具体操

avatar 2048 AI社区

解析关键要点！AI应用架构师通过智能体探索量子领域的思路

我们用Gym库定义一个量子算法优化的环境，就像“智能体的游戏地图”——智能体在这个环境里“调整参数”，获得“奖励”（能量越低，奖励越高）。import gym# 1. 初始化H₂分子的电子结构问题（目标哈密顿量）# 2. 动作空间：调整4个量子电路参数（每个参数的调整幅度±0.1π）# 3. 观测空间：当前的4个参数（范围0~2π）# 4. 量子模拟器（用经典计算机模拟量子实验）# 5. 初始化参

avatar 2048 AI社区

零售业AI Agent应用策略

随着人工智能技术的飞速发展，AI Agent在各个领域的应用日益广泛，零售业也不例外。本文章的目的在于深入探讨AI Agent在零售业中的应用策略，涵盖从顾客服务、商品推荐到库存管理等多个方面。通过对相关技术和实际案例的分析，为零售企业提供可操作的建议和策略，以提高运营效率、增强顾客体验和提升竞争力。研究范围包括传统实体零售和电商零售，旨在为不同类型的零售企业提供通用且具有针对性的AI Agent

avatar 2048 AI社区