复现GPT-2 124M这件事我前后折腾了大概三周从最初在单卡RTX 4090上跑通一个能生成连贯英文句子的模型到后来把loss曲线压到和官方原始权重基本对齐中间踩了不少坑也把整个Transformer的训练机制彻底摸了一遍。如果你也想从零复现GPT-2 124M这篇文章就是我整理出来的完整路线和实操记录会告诉你每一步怎么做、为什么这么做、以及最容易被卡住的几个地方在哪。先说明一下这里的“从零复现”不是从数学公式开始推导也不是从纸上徒手写CUDA核函数而是指不加载任何预训练权重从随机初始化出发自己搭建数据管线、实现GPT-2模型结构、跑通训练流程最终得到一个行为和原始GPT-2 124M基本一致的模型。这件事听起来唬人实际上用一块24GB显存的消费级显卡就能起步重点在于把架构细节和训练策略吃透。这个项目适合谁两类人一类是想真正理解Transformer内部机制、不想停留在“调库”层面的算法工程师和学生另一类是准备做LLM微调、RLHF、模型压缩等方向需要先打好底层功底的研究者。GPT-2 124M是最小但最完整的“大语言模型标本”它的每个设计选择都直接沿用了后来GPT-3、LLaMA等模型的核心思路你把它彻底搞明白后面看任何开源模型都能很快上手。1. 项目拆解与复现路线规划1.1 为什么选择GPT-2 124M作为切入点市面上能复现的模型很多小到Karpathy的miniGPT大到几百B参数的MOE模型但我最终还是锁定了GPT-2 124M。原因其实很朴素它足够简单又足够“真”。124M这个数字是参数量具体拆开来看12层Transformer解码器、每层12个注意力头、隐藏维度768、上下文长度1024词表大小50257。这个规模和今天动辄几十B的模型比起来确实小但它的架构不是玩具——自回归语言建模、因果自注意力、残差连接、LayerNorm、BPE分词现代LLM该有的机制它全都有。选它还有一个非常现实的好处OpenAI官方公开了124M的原始权重你复现完以后可以和官方的logits做对比。这一点太重要了我复现技术类项目最大的痛苦就是不知道自己做对了没有而GPT-2 124M给了你一个“标准答案”——如果你的模型和官方权重在相同输入下输出的top token分布高度相似说明你的代码、超参数、数据分布基本都对了。另外一个原因是成本可控。124M参数在FP16精度下权重只占约250MB训练时用4090级别的显卡就能跑通完整流程。我试过在同样一块卡上复现GPT-2 1.5B显存完全不够要做各种offload和梯度检查点复杂度一下子翻了几倍根本不适合作为“第一次用双手从零搭一个LLM”的入门项目。1.2 核心模块拆解数据、模型、训练、评估把整个项目拆开看其实是四条平行的线第一是数据管线。GPT-2官方用的是OpenWebText但这个数据集有40GB左右下载和清洗都比较费劲。我在第一版复现里用的是有人提前处理好的OpenWebText子集大约9GB的token序列足够训练出一个loss在3.2左右、能生成连贯文本的模型。如果你只是想把代码流程跑通用TinyStories这种更小的数据集也行后面我会单独讲。第二是模型架构。GPT-2的核心是Transformer解码器块12个Block堆叠。每个Block包含一个带因果掩码的多头自注意力子层和一个两层MLP子层两个子层前都有LayerNorm后都有残差连接。这部分我强烈建议自己手写一遍哪怕代码慢一点也比直接抄nanoGPT要强得多。第三是训练流程。预训练用的是标准的自回归语言建模目标优化器用AdamW学习率6e-4并配合warmup和余弦退火。这里面有很多细节比如梯度累积步数怎么算、梯度裁剪阈值设多少、Float16混合精度下怎么防止溢出这些都要自己调一遍才有体感。第四是评估环节。loss当然是最直接的指标但你还需要一个“模型真的学到了知识”的验证方式。我用了HellaSwag这个常识推理基准124M模型能达到官方公布的约29%左右准确率如果你能跑到这个数附近说明复现是成功的。1.3 我采用的总体技术路线我的整体路线分为四步每一步都有明确的“完成标志”第一步搭好基础环境并验证数据管线能产出正确的token序列。这一步的完成标志是你能从一个文本文件中采样出一批形状为[batch_size, block_size]的token张量并打印出解码后的文本确认没有问题。第二步实现GPT-2模型结构做一次前向传播。完成标志是输入一个batch的token序列模型能正常输出logits同时loss是一个有限的数值而不是NaN。第三步实现完整训练循环小规模试跑。用TinyStories数据跑几百步确认loss能稳定下降。第四步切换到OpenWebText完整数据启动正式训练按计划跑完20万步左右并用HellaSwag评估最终效果。整个过程中我使用的工具链是PyTorch 2.x CUDA 12 HuggingFace的tokenizers库 Weights Biases来做训练曲线可视化。为什么不直接用HuggingFace的GPT2Model因为复现的意义就在于自己把每一行代码写出来HF的模型实现适合做对照实验不适合当起点。2. GPT-2模型架构逐层拆解与自实现2.1 模型整体结构与关键类设计GPT-2 124M的参数量分布大概是这样的词嵌入矩阵加上输出投影矩阵占了最大头50257 x 768 x 2约77M剩下的是12个Block中的注意力层和MLP层。注意两个细节第一GPT-2的输入embedding和输出lm_head是共享权重的参数绑定Weight Tying这样既省了约38M参数又在一定程度上让模型学习更稳定第二GPT-2用的是可学习的位置编码不是后来GPT-Neo和LLaMA用的RoPE。我实现模型时的核心类设计如下import torch import torch.nn as nn import torch.nn.functional as F class LayerNorm(nn.Module): def __init__(self, dim, eps1e-5): super().__init__() self.eps eps self.gamma nn.Parameter(torch.ones(dim)) self.beta nn.Parameter(torch.zeros(dim)) def forward(self, x): mean x.mean(-1, keepdimTrue) var x.var(-1, keepdimTrue, unbiasedFalse) return self.gamma * (x - mean) / torch.sqrt(var self.eps) self.beta这里有个容易被坑的细节实现LayerNorm时var一定要用unbiasedFalse也就是不做贝塞尔校正。PyTorch内置的nn.LayerNorm默认就是这么算的但如果你自己手写并且很自然地用了偏估计最终loss会偏高一点点而且很难发现问题。因为偏差不大看起来“模型也是能跑的”但复现的精度对不上。然后是核心的CausalSelfAttention、MLP和Block类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) self.c_proj nn.Linear(config.n_embd, config.n_embd) self.n_head config.n_head self.n_embd config.n_embd self.register_buffer(bias, torch.tril(torch.ones(config.block_size, config.block_size)) .view(1, 1, config.block_size, config.block_size)) def forward(self, x): B, T, C x.size() qkv self.c_attn(x) q, k, v qkv.split(self.n_embd, dim2) k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) q q.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) att (q k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) att att.masked_fill(self.bias[:, :, :T, :T] 0, float(-inf)) att F.softmax(att, dim-1) y att v y y.transpose(1, 2).contiguous().view(B, T, C) y self.c_proj(y) return y这段代码里有一个关键的参数初始化细节c_proj层的权重在GPT-2里是按1/sqrt(2*n_layer)缩放的这个不是随便加的而是后面训练稳定的重要设计。我一开始没做这个缩放训练时loss确实能下降但偶尔会出现loss尖峰训练极不稳定。后来查阅GPT-2的源码才发现这个操作加上之后训练平滑了很多。2.2 因果注意力掩码的实现细节因果掩码是自回归模型的核心机制它的含义是第T个位置的token只能看到前T-1个位置的信息不能偷看未来的内容。实现方式有两种我强烈建议用上面这段代码里的方式——用一个预先计算好的下三角矩阵在softmax之前把未来的位置mask成负无穷。为什么不直接用attention_mask张量参与矩阵乘法呢因为在PyTorch里masked_fill softmax的组合在CUDA上有专门的融合优化速度快很多。而掩码矩阵注册为bufferregister_buffer而不是普通属性是为了让它在模型迁移到GPU、切换dtype时自动跟随。另一个实现选择是在计算注意力分数时除以sqrt(d_k)。这个缩放因子在原始Transformer论文里是常数1/sqrt(d_k)它的作用是防止点积结果过大导致softmax进入饱和区。如果不做这个缩放模型的训练会非常不稳定尤其是在FP16混合精度下很容易出现NaN。2.3 LayerNorm位置从Post-Norm到Pre-NormGPT-2和原始Transformer有一个非常关键的区别原始Transformer用的是Post-Norm即每个子层的顺序是“注意力/MLP → 残差连接 → LayerNorm”而GPT-2用的是Pre-Norm顺序是“LayerNorm → 注意力/MLP → 残差连接”。改这个位置的动机在后来的研究里已经被反复验证Post-Norm在深层网络中梯度容易消失需要额外做warmup和复杂的初始化技巧Pre-Norm则天然更稳定因为残差连接的路径是干净的梯度可以直接从输出层回传到输入层。但Pre-Norm也有一个副作用模型最终的输出层之前没有额外的normalization所以GPT-2在最后的lm_head之前还需要补一个最终的LayerNorm。这个细节很多人会漏掉漏掉之后模型仍然能训练但最终的logits分布会偏大困惑度Perplexity会明显变差。我把GPT-2 Block的实现写出来以后用一句话总结它的结构每一个Block 残差连接包裹的注意力子层 残差连接包裹的MLP子层每个子层内部先做LayerNorm再做计算。这个“块”的模板后来被几乎所有开源大模型沿用包括LLaMA系列。你把这个Block彻底理解了看任何现代LLM源码都会轻松很多。2.4 激活函数、嵌入层与权重初始化GPT-2的MLP中间层用的激活函数是GELU而不是ReLU。GELU和ReLU的关键区别在于它在负数部分不是完全截断的而是有一个平滑的非线性过渡。理论上GELU能提供更平滑的梯度流实践中也确实会在困惑度上稍微好一点。实现GELU有一个引战级的问题用精确的误差函数版本还是用tanh近似版本。GPT-2原始代码用的是tanh近似def gelu(x): return 0.5 * x * (1.0 torch.tanh(math.sqrt(2.0 / math.pi) * (x 0.044715 * torch.pow(x, 3))))PyTorch内置的nn.GELU默认用的是精确版本。两个版本数值上差异极小但如果你做“复现”并且要和官方logits对比建议使用tanh近似版本不然有些激活值在最后几位对不上。权重初始化是另一个容易忽视但影响极大的细节。GPT-2的初始化策略是所有Linear层的权重从均值为0、标准差为0.02的正态分布中采样LayerNorm的gamma初始化为1、beta初始化为0而每个Block中c_proj层的权重标准差要额外除以sqrt(2n_layer)。这个除以sqrt(2n_layer)是GPT-2引入的一个创新点目的是让每个残差分支的贡献在深层网络中不至于累积过大从而保证深层模型的稳定性。我在复现时没有引入这个初始化结果在训练到几千步的时候loss曲线出现了周期性的尖峰模型几乎要发散。后来查了GPT-2的源码才发现这个细节加上之后几十万步训练再也没出过问题。3. Tokenizer与数据管线从文本到Token序列3.1 字节级BPE分词器原理解析GPT-2使用的分词器是字节级BPEByte Pair Encoding的一个变体。传统的BPE是自己维护一份词表把单词拆成词片而字节级BPE不是以字符为最小单位而是以UTF-8的字节为最小单位。这样做的好处是无论什么语言、什么特殊符号都能被字节级编码覆盖不存在“OOV”超出词表的问题。它的训练过程大致是先从语料中统计所有字节的出现频率然后反复合并出现频率最高的字节对直到词表大小达到目标GPT-2是50257。这一合并过程会留下一个merge表推理时按同样的规则把文本逐步合并成词表中的token。一个具体的例子单词“hello”在字节级BPE中可能会被拆成“h”、“e”、“ll”、“o”这样的token组合。注意“ll”是一个合并后的词片因为这个组合在训练语料里出现的频率足够高。理解BPE的细节很重要因为它是模型“看到”文本的方式。如果你用了一个分词方式和训练数据不一致的tokenizer相当于把模型的输入空间完全搞乱了模型再强也没用。我在初学阶段就做过这种蠢事用GPT-2的tokenizer处理数据但词表加载错了版本用成了GPT-2 Medium的词表结果生成的内容全是乱码排查了半天才发现是tokenizer的词表大小和模型配置对不上。3.2 从零实现BPE还是直接用tiktoken这里我的建议很明确如果目标是“复现GPT-2并训练出好模型”直接用OpenAI开源的tiktoken库不要自己重新实现BPE。原因是BPE的合并规则实现起来代码量不小而且极容易在边际情况比如空白字符、多语言文本、未知字节上出错。第一天我把目标定在“从零实现GPT-2”当时也雄心勃勃地想手写tokenizer结果浪费了两天时间在调试各种UTF-8边界情况上最后换回tiktoken才把主线跑通。但要注意tiktoken的对应编码名称是gpt2加载的时候要确认编码名称否则你会拿CLIP或者别的模型的tokenizer来凑导致词表不匹配import tiktoken enc tiktoken.get_encoding(gpt2) encoded enc.encode(Hello, world!) print(encoded) # 输出类似 [15496, 11, 995, 0]tiktoken会额外处理一个细节特殊token。比如GPT-2的词表中保留了|endoftext|作为文档分隔符tiktoken在encode时默认不会输出这个token但你可以显式指定allowed_special参数来控制它。在预训练时我建议每段文档之间插入一个|endoftext|帮助模型学会“文档边界”这个概念。3.3 数据加载与批处理的关键细节数据管线的核心任务是把原始文本文件转成训练的token序列。这一步比想象中更讲究我最初的做法是把所有文本拼接成一个巨大的一维token数组然后用一个偏移量随机抽取block_size长度的片段作为训练样本。这个做法的好处是简单而且能最大化利用数据坏处是模型会看到跨越文档边界的token片段产生噪声。GPT-2官方处理这个问题的方式是插入|endoftext|作为文档分隔符然后在切分数据时随机偏移。具体逻辑是对每个文档tokenize之后在后面加上一个|endoftext|的token id然后把所有文档拼接成大数组。读取训练样本时不是从固定位置切而是在每个block的起始位置加一个0到block_size-1之间的随机偏移这样每次epoch的切分位置都不同相当于给训练数据引入了随机性。数据加载器我用的是PyTorch的DataLoader加自定义Dataset但有一个更高效的方案把token数组保存为numpy的uint16格式然后在Dataset里用np.memmap做内存映射。这样即使数据是几十GB也不需要全部读入内存训练时按需读取速度和内存占用都很理想。import numpy as np import torch from torch.utils.data import Dataset class TokenDataset(Dataset): def __init__(self, data_path, block_size): self.data np.memmap(data_path, dtypenp.uint16, moder) self.block_size block_size def __len__(self): return len(self.data) - self.block_size def __getitem__(self, idx): shift torch.randint(0, self.block_size, (1,)).item() x torch.from_numpy(self.data[idx shift : idx shift self.block_size].astype(np.int64)) y torch.from_numpy(self.data[idx shift 1 : idx shift self.block_size 1].astype(np.int64)) return x, y这段代码里有两个容易被忽略的地方。第一为什么要转成np.int64因为模型的embedding层接受的索引必须是int64uint16会直接报错。第二x和y的关系是“下一个token预测”y是x的输入序列整体右移一位的结果训练目标是给定前T个token预测第T1个token。3.4 词表大小、上下文长度与显存的数学关系在开始训练之前我建议你心里先有一个显存占用的估算。假设batch_size为1、上下文长度为1024、词表大小为50257。前向传播时logits的形状是[1, 1024, 50257]在FP32下约占205MB在FP16下占约103MB。然后loss函数要在词表维度上做softmax和交叉熵中间还要产生概率分布峰值显存可能翻倍。124M模型的实际显存占用大头在AdamW优化器的状态上。AdamW要为每个参数保存一阶动量和二阶动量各占4字节FP32124M参数就是1GB左右。加上模型权重本身约500MBFP32再算上梯度、激活值、logits综合下来一块24GB显存的卡是够用的但也不是非常宽裕。如果你只有12GB甚至8GB显存也不需要绝望。两个思路一是用梯度累积把batch拆成micro batch每个micro batch前向和反向计算完就释放中间张量二是用DeepSpeed的ZeRO-Offload或者CPU offload技术不过这就属于“增加学习难度”的范畴了我建议先跑通全流程再考虑优化。4. 训练全流程超参数、策略与loss曲线观察4.1 优化器选型与超参数详解GPT-2复现的超参数基本是公开的但每一项背后的逻辑值得说清楚。优化器用AdamW而不是Adam区别在于权重衰减weight decay的实现方式。Adam把权重衰减以L2正则的形式加在梯度上而AdamW把权重衰减直接从更新公式中减掉。理论和实践都表明AdamW在Transformer训练中表现更好。PyTorch里直接用torch.optim.AdamW即可。学习率我使用的是6e-4这是GPT-3论文给出的Base设置适用于124M~1.5B左右的模型。这个数字不是拍脑袋定的过大会导致训练初期非常不稳定过小会导致收敛速度太慢。如果你用TinyStories等小数据集做调试建议把学习率降到3e-4左右因为数据量少、分布简单用小学习率能让loss下降得更平滑。权重衰减设为0.1但有一个重要细节只有参与乘法计算的权重矩阵才做权重衰减LayerNorm的gamma/beta和bias不应做衰减。实现时用PyTorch的param_group把参数分两组def configure_optimizers(model, weight_decay, learning_rate, betas, device_type): param_dict {pn: p for pn, p in model.named_parameters()} 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} ] optimizer torch.optim.AdamW(optim_groups, lrlearning_rate, betasbetas) return optimizerdim大于等于2的通常都是权重矩阵bias和LayerNorm参数都是1维这样区分能保证权重衰减只作用于真正的权重。4.2 学习率调度warmup与余弦退火学习率调度是训练稳定性的关键环节。GPT-2/3的官方实现使用了一个组合策略先线性warmup再余弦退火到最终学习率。warmup的步数是716步这个数字是怎么来的实际上是总训练步数大约20万步的0.3%左右GPT-3论文里给出的经验值是总步数的0.3%。你可以这样理解训练刚开始时模型权重是随机初始化的梯度的方向和尺度都不稳定如果直接使用6e-4这种较大的学习率很容易导致参数越过最优区域甚至发散。warmup阶段让学习率从0线性增长到6e-4相当于给模型一个“预热”的过程让梯度方向逐渐稳定下来。余弦退火则是让学习率在整个训练后半段以一个余弦曲线平滑下降到接近0的值。这样做的好处是训练后期可以使用较小的学习率精细调整参数避免在loss最小值附近震荡却下不去。我在实现时用的调度函数如下def get_lr(it): if it warmup_iters: return lr * it / warmup_iters it min(it, max_iters) decay_ratio (it - warmup_iters) / (max_iters - warmup_iters) coeff 0.5 * (1.0 math.cos(math.pi * decay_ratio)) return final_lr coeff * (lr - final_lr)final_lr我设成了6e-5即初始学习率的10%。这个比例是经验值最终学习率不完全归零的原因是为了让模型最后不会在优化面上卡死保留一点活动空间。4.3 批大小、梯度累积与混合精度GPT-2官方训练的batch大小是512个序列每个序列1024个token换算成tokens就是约52万个token。在现代家用显卡上一次前向和反向不可能放下512个样本所以梯度累积是必须的scaler torch.cuda.amp.GradScaler() optimizer.zero_grad(set_to_noneTrue) for micro_step in range(gradient_accumulation_steps): logits, loss model(X_batch, Y_batch) loss loss / gradient_accumulation_steps scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update()注意loss除以梯度累积步数这一步否则每个micro batch的loss都会独立地贡献梯度等效学习率会被放大gradient_accumulation_steps倍。混合精度方面我使用torch.cuda.amp自动混合精度即大部分计算在FP16下完成但优化器更新和梯度裁剪在FP32下进行。FP16能显著降低显存占用并提升训练速度但会引入一个隐患梯度值如果太小会直接变成0导致训练卡住。GradScaler的作用就是动态调整损失缩放系数避免下溢出。梯度裁剪的阈值我设为1.0。这个值的意思是如果梯度的全局L2范数超过1.0就把所有梯度按比例缩小到1.0。它防止的是个别极端梯度把参数冲出稳定区间。在训练初期梯度范数有时会在5以上裁剪后训练稳定很多。4.4 训练过程实测loss曲线与生成效果演进我在4090上正式训练用的超参数如下这是完整跑完20万步的配置超参数数值模型参数124M上下文长度1024Batch大小等效512个序列 / 约52万token梯度累积16步 × 32个micro batch学习率6e-4最终学习率6e-5Warmup步数716最大训练步数200000Adam betas(0.9, 0.95)权重衰减0.1梯度裁剪1.0训练初期0~700步loss从10.9快速下降到6.0左右这其实是模型在快速学习词法层面的规律比如字母组合、常见单词的拼写。到几千步后loss降速明显放缓进入一个平台期这时模型开始学习语法结构和短语搭配。一个值得记录的观察点是训练到约1万步时模型生成的文本虽然语法上基本正确但语义上完全是胡言乱语。到5万步左右能生成连贯的中等长度句子但会频繁重复同一句话。到12万步以后模型的重复问题明显缓解可以生成结构完整的小短文虽然内容依然常常荒谬但句与句之间的相关性明显增强。最终训练的loss在验证集上约为3.18换算成Perplexity大约是24。这个和官方的2.97左右还有一些差距主要原因是数据集规模不足——官方用了40GB的OpenWebText我只用了9GB的子集。5. 常见问题与排查技巧实录5.1 典型故障速查表下面是训练过程中我遇到的最常见的五个问题以及对应的排查思路问题现象根本原因解决方案loss一开始就是NaN学习率过大、FP16下数据溢出、初始化有问题检查是否是输出层scale问题把混合精度关闭先跑FP32确认把学习率降到1e-5看是否恢复正常loss停在某个值不下降数据管线有bug比如label等于输入模型在“作弊式学习”打印一个batch的x和y确认y确实是x右移一位训练时显存不足OOMmicro batch过大或梯度累积方式不对把micro batch减半调高梯度累积步数检查是否需要gradient checkpointing验证集loss下降但生成效果很差模型过拟合到训练数据或生成时未做采样训练时loss本身是困惑度的指标生成时要用top-k/top-p采样不能直接取argmax复现结果和官方checkpoint差距较大权重初始化方法不同、数据分布差异、学习率调度细节不同逐一对比先固定随机种子再对比官方源码的初始化细节最后评估用HellaSwag而非loss原始值5.2 如何在固定Step观察模型“是否真的在学”光看loss曲线还不够我强烈建议你每隔一定步数做三件事。第一件生成一段文本。从固定的prompt开始比如“In the beginning”采样一段长度100左右的文本直接观察语义的变化。这是最直观的“模型进步感”来源。第二件对比验证集loss和训练集loss的差距。如果训练集loss持续下降但验证集loss几乎不动甚至回升说明模型开始过拟合此时需要考虑增加数据量、增加dropout或者提前停止。第三件记录一个固定样本的top-k预测。选一个测试样本在训练的不同阶段打印模型预测概率最高的前5个token。你会看到模型从开始预测完全无关的词到后来能预测出显而易见的词这个变化对理解模型内部能力的发展非常有帮助。5.3 利用官方检查点验证复现正确性我复现完成后做了一次非常关键的验证把官方发布的GPT-2 124M权重加载到我的代码框架里然后对相同输入跑前向对比我训练的模型和官方模型输出的token概率分布。具体做法是from transformers import GPT2LMHeadModel model_hf GPT2LMHeadModel.from_pretrained(gpt2) my_state_dict model.state_dict() hf_state_dict model_hf.state_dict() # 打印所有键名逐层对比参数名称和形状 for k in my_state_dict: if k in hf_state_dict: print(k, my_state_dict[k].shape, hf_state_dict[k].shape) else: print(missing in HF: , k)刚加载时只需要解决参数名的映射问题因为我的命名和HF不完全一致但形状肯定要完全匹配。修改完以后对同一个输入文本我模型和官方模型的last layer hidden state余弦相似度到了0.99以上这说明架构实现完全正确。这一步极其推荐你去做它比你调多少训练超参数都能更快让你确认自己代码有没有问题。如果架构对不上后面训练再久也是白费。5.4 一些小众但实用的避坑技巧第一训练过程中如果发现loss出现周期性的尖峰大概率不是数据问题而是权重初始化时c_proj层的缩放没有做。这个之前提过再强调一次。第二不要在一开始就用FP16混合精度。先用FP32跑通200步确认模型能正常下降再打开混合精度。否则你会在“代码bug”和“精度溢出”两个问题之间来回排查非常痛苦。第三如果训练速度太慢不是显卡不够好而可能是数据加载成为瓶颈。用memmap和DataLoader的num_workers可以解决大部分问题。第四显存不够时除了梯度累积还有一个技巧是使用梯度检查点gradient checkpointing。PyTorch里调用model.gradient_checkpointing_enable()即可。它用“反向传播时重新计算前向结果”来换取显存代价是训练速度下降约30%但显存占用可以降低一半以上。第五如果机器重启后训练中断要恢复训练必须保存完整的训练状态不只是模型权重。我建议每个checkpoint至少包含model_state_dict、optimizer_state_dict、scaler_state_dict、当前步数、当前学习率。少了scaler状态下一次恢复训练时混合精度的损失缩放因子会重置有时会导致短期训练不稳定。6. 从“复现成功”到“更进一步”可扩展的方向复现完成之后GPT-2 124M不仅仅是一个学习项目它完全可以继续演变成更深入研究的起点。第一个方向基于你自己的权重做监督微调SFT。用几百条手工标注的“问题-回答”对在预训练权重上做几轮微调就能得到一个简单的对话模型。这个过程本身又是一个全新的学习项目你能切身体会到指令微调如何“激活”预训练模型已有的知识。第二个方向研究注意力机制和可解释性。124M模型足够小你完全可以单独提取出某个注意力头在特定层的行为看看它对句子的什么特征最敏感。比如第8层的某个注意力头可能对句子的主谓关系敏感第10层的一个头可能更关注语义相关的词。这种研究在更大的模型上反而不容易做。第三个方向做结构化剪枝和蒸馏。因为有完整的训练代码你可以尝试剪掉一些不重要的注意力头观察性能下降情况或者训练一个更小的模型去蒸馏124M模型的行为。这些都是实际工程中常用的技术在这样一个中等模型上做实验成本可控。第四个方向把训练框架改成多卡分布式。我全程是在单卡上完成的但如果把数据并行加进来用DistributedDataParallel再跑一遍就相当于把项目延伸到了“大规模训练系统”领域。这一步的代码改动其实不大但对理解分布式训练中的通信开销、batch划分、梯度同步非常有帮助。我个人在实际操作中最大的体会是从零复现一个模型最大的收获不是那个124M权重文件本身而是你被迫把每一个细节——从tokenizer的合并规则到AdamW的参数分组——都彻底弄明白了。这些细节在平时调库的时候根本意识不到一旦自己动手它们全是决定成败的关键。如果你也想真正理解大模型是怎么训练出来的找一个下午打开编辑器从导入PyTorch开始一行一行写出自己的GPT-2那种感觉和看一百篇教程都不一样。最后再分享一个小技巧训练一个大型模型需要耐心loss不降or生成效果差的时候不要马上怀疑代码逻辑先检查数据管线的输出是否符合预期再检查初始化最后才去翻训练超参数。这几乎是所有训练Debug的通用顺序。希望这篇记录能帮你少走一些弯路。
