← 返回首页
07_pretraining_viz.html

真训练:把玩具 GPT 喂真实数据(预训练 GPT-2 124M)

本页对应 phase1-124m/04_gpt2_124m.py。架构内核和前几关一模一样(注意力 + FFN + 残差 + LayerNorm), 变的是规模和一整套"让大模型训得稳、训得动"的训练工程。这一页把"配置怎么涨 → 真数据怎么喂 → 显存不够怎么办 → 学习率怎么排 → 最后收敛到哪"五步拆开,每步亲手拨一个旋钮。

STEP 1
从玩具到真模型
STEP 2
真数据 + BPE + shards
STEP 3
梯度累积
STEP 4
学习率 schedule
STEP 5
收敛 + 数据量对照
GPT-2 124M n_layer 12 n_head 12 n_embd 768 block_size 1024 vocab 50304 total_batch 524288 token/步 max_lr 6e-4 这套就是 GPTConfig 真值 · 配套 04_gpt2_124m.py

① 从玩具到真模型:同一套积木,数字全部放大

第 3 关那个字符级小 GPT(nanoGPT 风格)和这一关的 GPT-2 124M,零件清单完全一样 —— 区别只是每个数字往上跳了一大截。点下面两个按钮,看同一张配置表里哪一列被点亮,以及每行涨了多少倍。

超参字符级玩具(03)GPT-2 124M(04)
vocab 为什么是 50304,不是 GPT-2 真实的 50257?
GPT-2 的 BPE 词表真实大小是 50257。 代码里把它向上取整到 128 的倍数 → 50304(GPTConfig.vocab_size)。原因纯粹是硬件: GPU 的矩阵乘对"维度是 2 的幂 / 128 的倍数"算得更快。多出来的 50304-50257=47 行 永远不会作为真实 token 出现,学不到也无害,白拿一点速度。
玩具是"字符级",GPT-2 是"BPE",差在哪?
第 3 关 vocab 只有 65vocab —— 那是把莎士比亚文本里出现的不同字符(字母、标点、空格)去重数出来的, 模型一次只预测一个字母。GPT-2 用 BPE,一个 token 往往是一截子词(如 " the""ing"), 50304 个 token 覆盖了真实英文,所以同样长度能塞下多得多的信息。架构没变,只是"最底层的字典"换大了。

模型壳子放大了,但真模型要喂真数据。下一步:FineWeb-Edu 语料怎么变成能一块块喂进去的东西。

② 真数据 + 真 BPE + shards:把 10B token 切成块,顺序喂

训练数据是 FineWeb-Edu: 先用 GPT-2 的 BPE 把文本切成 token id,再把上百亿个 id 存成一批 shard (.npy 文件)。训练时 DataLoaderLite 一片片顺序读、读完换下一片。点「▶ 开始喂数据」看指针滑过。

FineWeb-Edu 原文 GPT-2 BPE 分词 token id 存成 .npy shards 顺序喂训练
BPE 示意一句话被切成 token(下面切法是示意,非真实 BPE)
shards10B token 切成多片,指针顺序滑过
每片是一个 .npy(int token id);DataLoaderLite.next_batch 切出 B×T+1 个连续 token,前 B×T 个做输入、后移一位做标签。读完一片换下一片。
为什么不一次性把整个数据集读进内存?
10B token 即使按 2 字节存也有几十 GB,塞不进显存、也未必塞得进内存。切成 shard 后,DataLoaderLite 同一时刻只 np.load 一片到内存,指针 pos 在片内滑动, 到头了 cur_shard+1 换下一片(循环)。这样无论数据多大,内存占用是常数
上面这串 token 切法是真的吗?
不是,上面只是示意。真实 GPT-2 BPE 词表有 50304 行,切分规则复杂(按字节对频率合并), 一个空格、一个词尾都可能单独成 token。这里只为让你直观看到"一句话 → 若干 token"这件事, 别把具体切法当真。真实分词在 prepare_fineweb.py 里用 GPT-2 的 tokenizer 完成。

数据备好了,但 GPT-2 一步要吃 524288 个 token —— 单卡一次根本喂不下。下一步:梯度累积

③ 梯度累积:显存放不下大 batch,就攒梯度

GPT-2 的等效大 batch 是 524288token/步(total_batch_size,219)。 单卡一次只喂得下一个微 batch:8micro_batch 条序列 × 1024seq_len = 8192 token。 办法:连喂若干个微 batch、把梯度累加进"蓄水槽",攒满了才 optimizer.step() 走一步。拖滑块看要攒几次。

64
1(不累积)64 = 真实默认128
放水线 = step()
524,288 token / 优化步(等效 batch)
累积梯度,为什么 loss 要除以累积次数?
因为大 batch 的目标是对 524288 个 token 的损失求平均。代码里 loss = loss / grad_accumloss.backward():每个微 batch 贡献 1/N, 累加 N 次正好等于"在一个大 batch 上算一次平均梯度"。不除的话梯度会被放大 N 倍,等于偷偷把学习率乘了 N。
累积出来的大 batch,和"真的一次喂这么大"完全等价吗?
数学梯度上等价(忽略累加的浮点误差)。差别只在含 batch 统计量的层 —— 好在 GPT-2 用的是 LayerNorm(只在单个 token 内部归一化,不跨样本), 所以累积和真大 batch 在这里就是等价的,可以放心拿显存换 batch。

大 batch 攒好了,每一步该用多大的学习率?大模型不能一上来就猛冲。下一步:warmup + 余弦退火。

④ 学习率 schedule:先热身,再余弦退火(本页交互重点)

学习率不是常数。前 warmup_steps热身 步从 0 线性爬到峰值 6e-4max_lr,之后余弦平滑退火到 6e-5min_lr(= max_lr × 0.1),末尾保持不变。 拖两个滑块,曲线实时精确重画(就是 get_lr 这个纯函数,不是示意)。

715
0(不热身)715 = 10B 跑3000
19073
200019073 = 10B 跑25000
学习率 lr(step) warmup 区 峰值 max_lr=6e-4 谷底 min_lr=6e-5
为什么大模型不 warmup、开局就崩?
刚初始化时参数是随机的,梯度又大又乱。这时若直接上峰值学习率,一步就能把参数推到很离谱的地方, loss 直接发散(NaN)。warmup 让学习率从 0 慢慢爬,给模型几百步"找到合理方向"的时间, 等梯度稳定了再加速。get_lr 里就是 step < warmup_stepsmax_lr × (step+1) / warmup_steps 这条线性斜坡。
为什么后面要余弦退火,不一直用大学习率?
大学习率适合"大步快走"快速下降,但接近谷底时步子太大会在最优点附近反复横跳下不去。 余弦退火让学习率越走越小,后期小步精修、稳稳收敛。代码里是 coeff = 0.5×(1+cos(π·ratio)),把进度 ratio 从 0→1 平滑映射成系数 1→0, 再插值到 [min_lr, max_lr] 之间。
优化器那一行还有哪些"老司机"设置?
优化器是 AdamW: betas=(0.9, 0.95)(GPT 系比默认 0.999 更小的二阶动量)、eps=1e-8fused=True(把更新融合成一个 CUDA kernel 提速)。 weight_decay=0.1分组:只对 2 维参数(矩阵、嵌入)做衰减,bias 和 LayerNorm 的增益/偏置不衰减 (configure_optimizers 里的 decay / nodecay 两组)。另外每步还做梯度裁剪把范数截到 1.0,防偶发巨梯度炸训练。

学习率排好了,训练真跑起来 —— loss 会从哪降到哪?下一步用两个真实端点看收敛和数据量的关系。

⑤ 收敛 + 数据量对照:两个真实端点,数据多 33×

随机初始化时,模型对 50257 个 token 完全瞎猜,loss ≈ ln(50257) ≈ 10.82(代码开训前会打印这个基线)。 训练让它一路下降。本项目没有保存逐 step 的 loss,只有两次跑各自的终点 —— 下面是这两个真实端点, 以及它们对应的数据量差 33×

短跑 · 约 300M token
3.65 val_loss
step 1999 · 快速验证管线
(用了比 524288 更小的 batch)
长跑 · 约 10B token
3.02 val_loss
step 19072 · max_steps=19073
warmup 715 · 完整余弦退火
⚠️ 示意下降趋势(非实测逐 step,只有两端点是真的)
真实端点 1:300M → 3.65 真实端点 2:10B → 3.02 示意趋势(非实测)
数据量从 300M → 10B 翻了约 33×,val_loss 从 3.65 → 3.02。 虚线只是示意大致形状,中间没有任何实测点 —— 别把它当真实曲线。两个圆点才是项目里真实跑出来的数。
那 bf16 / torch.compile / Flash-Attention 是什么,这页为什么不细讲?
它们是纯工程提速,不改训练数学:bf16 混合精度(算得快、省显存)、 torch.compile(把模型编译成更快的 kernel)、Flash-Attention (F.scaled_dot_product_attention,不显式建 T×T 矩阵)。它们让同样的训练更快更省,但 loss 曲线该到哪还是到哪 —— 所以本页只点到为止。
loss 3.02 到底算好还是不好?
从 10.82 的随机基线降到 3.02,已经是一个能写出通顺英文的 base 模型了(GPT-2 124M 原版量级)。 但它仍只会续写、不会答题。而且 loss 还没到底:继续喂更多 token、训更多步还能再降 —— 这正是"数据量 33× → loss 明显更低"想说的:规模和数据,是预训练最朴素也最硬的杠杆

训练好了,这台机器怎么把字吐出来?去看「采样」那一关(05_sample.py / 05_sampling_viz)。点右下角「完成 🎉」。