首页 / 资讯中心 / 文章详情

迷你GPT训练全解:从Token化到训练循环的完整实践指南

迷你GPT训练全解:从Token化到训练循环的完整实践指南 ★ FEATURED ARTICLE
这一篇接着上一篇往下写。模型骨架搭完之后我对着空转的前向函数发了一会儿呆输入输出都定义了但它还不会说“人话”因为权重是随机的模型根本没有“见过”任何文本。要让一个迷你GPT真正开始“学习”要做的其实只有两件事把文本变成模型能消化的数据把梯度变成模型能吸收的营养。说白了就是数据工程和训练循环这两个环节。这篇文章就是来拆解这两个环节的。我会从文本怎么切Token开始讲再讲到怎么构造训练样本、怎么用DataLoader喂数据然后钻进训练循环内部把前向传播、损失计算、反向传播、优化器更新、梯度累积、学习率调度全部过一遍。最后附上我自己踩坑总结的问题速查表。适合正在手写GPT复现项目、或者刚看完Transformer结构但不知道数据怎么送进模型的读者。1. 数据工程文本是怎么变成“考题”的先聊数据工程。模型本身不识字只认识数字。你给它的每一段文本最终都要转成一串整数id再切成一条条“已知上文预测下一个字”的样本。如果样本切得不对、顺序打得不乱、验证集划分不合理后面训练循环写得再漂亮也没用。1.1 先把文本切成 TokenTokenization 是最基础的一步。人类语言是字符串模型只能吃数字向量所以需要一个词表把“字符”或“词”映射成整数id。对于学习原理来说我强烈建议先用字符级 tokenizer不做任何第三方分词直接对字符建表。原因很简单分词工具内部有太多隐藏逻辑出问题时不好判断是你的模型bug还是分词器的锅。字符级方案一眼看穿。代码其实就这几行def build_tokenizer(text): chars sorted(list(set(text))) stoi {ch: i for i, ch in enumerate(chars)} itos {i: ch for i, ch in enumerate(chars)} vocab_size len(chars) return stoi, itos, vocab_size def encode(s, stoi): return [stoi[c] for c in s] def decode(ids, itos): return .join(itos[i] for i in ids)注意一点词表大小直接决定embedding矩阵的尺寸也就是vocab_size * d_model词表太大会让模型参数膨胀太小又会让单条样本的序列太长。字符级词表通常几百个token就够训练速度飞快适合跑通整套流程。真实项目里更常用子词级别的方案但这里的核心目标是理解GPT的数据流先用字符级把整条链路打通后面要换分词器只是换个encode函数的问题。建完词表整个语料就可以变成一长串id。比如“人工智能正在改变世界”会变成[人, 工, 智, 能, 正, 在, 改, 变, 世, 界]这样的数字列表。到这里数据工程的第一步完成。1.2 构造样本上下文窗口与“预测下一个字”GPT的训练目标通俗讲就是“根据前面的内容猜下一个token”。所以一条样本由两部分组成输入是连续的一段token标签是把这段token往右平移一位的下一个token。假设原始序列是[a, b, c, d, e]上下文窗口block_size3那么样本应该是输入[a, b, c]标签[b, c, d]输入[b, c, d]标签[c, d, e]也就是说模型在每个位置都要输出“下一个token是谁”的预测标签就是它真正看到的下一个token。这种滑动窗口切法会让相邻样本大量重叠但这没关系。GPT训练本来就不在乎样本之间是否独立它需要的是让模型在海量“接龙”任务里总结统计规律。构造样本的代码很直接data encode(whole_text, stoi) block_size 128 # 第 i 个样本 x data[i : i block_size] y data[i 1 : i block_size 1]这里有个关键点跟前向传播的结构是对应的模型输入x得到(B, T, V)的logits第t个位置的logits去预测第t1个token所以标签和输入的长度是一样的只是错开了一位。很多新手在这一步会搞错把标签直接写成和x同一段那样模型等于在预测自己loss会降得特别快但完全没有意义。1.3 数据加载器与批量采样有了样本切分逻辑下一步就是要能高效地一批一批喂给模型。在PyTorch里最自然的做法是写一个Dataset类然后交给DataLoader去处理batch、打乱、多进程加载。我这里给一个比较实用的写法from torch.utils.data import Dataset, DataLoader import torch class TextDataset(Dataset): def __init__(self, data, block_size, random_sliceFalse): self.data data self.block_size block_size self.random_slice random_slice def __len__(self): return max(1, (len(self.data) - 1) // self.block_size) def __getitem__(self, idx): if self.random_slice: start torch.randint(0, len(self.data) - self.block_size, ()).item() else: start idx * self.block_size x torch.tensor(self.data[start:start self.block_size], dtypetorch.long) y torch.tensor(self.data[start 1:start 1 self.block_size], dtypetorch.long) return x, y关于random_slice这个参数我解释一下。训练集里我建议打开它也就是每个step随机从语料里取一个起点这样每次拿到的样本都不完全一样等于天然做了数据增强。验证集里必须关掉它保证每次评估用的都是同样的切片否则两个epoch的valid loss不在同一批数据上根本没有可比性。DataLoader配置示例train_loader DataLoader( train_dataset, batch_size8, shuffleTrue, num_workers0, drop_lastFalse, )drop_last我一般设为False因为后面梯度累积要处理不整除的情况。还有一点block_size就是模型的上下文长度它既要和模型内部的max_seq_len匹配也要考虑GPU显存。调大它能让模型看到更长上下文但显存占用会跟着涨尤其是注意力矩阵那块。1.4 训练集/验证集的划分与随机性数据工程里最容易忽视但又特别重要的一件事就是验证集的划分。我见过不少项目只做了一件事把loss打印出来然后看着训练loss下降就以为模型在学。但训练loss下降是必然的哪怕模型已经在死记硬背它也会降。想判断泛化能力必须有一个固定的验证集。常规做法是把整段语料按位置切分前90%训练后10%验证。注意验证集不能从中间随机抽因为你用的是滑动窗口样本随机抽会让一部分训练样本和验证样本内容重叠等于开卷考试。还要提一嘴随机种子。加载训练集数据的时候如果完全不设随机种子两次运行生成的训练样本不同这会让你很难判断“这次loss下降了到底是代码改好了还是运气好”。所以在准备数据之前我会先固定全局种子后面单独讲。2. 训练循环一次“学习”到底发生了什么数据准备好了接下来就是核心中的核心训练循环。很多人把它理解成“写一个for循环调一下loss.backward()就完事”确实是这样但要真正明白GPT是怎么学的就得拆开每一层看。2.1 前向传播模型先做一次“闭卷预测”每个step模型拿到一批输入x形状是(B, T)B是batch sizeT是block_size。经过embedding、多头注意力、前馈网络之后输出形状是(B, T, V)的logitsV是词表大小。每个位置上的logits可以理解成模型对这个位置“下一个token”的原始打分哪个token得分高模型就倾向于预测谁。此时模型还处于随机权重状态打分基本是乱来的但没关系因为后面有loss来纠正它。前向传播本身没有魔法它就是一次计算。但你要有意识地把“训练模式”和“推理模式”区分开。训练循环里要在模型前面加上model.train()因为像dropout这样的层在训练和推理时行为不一样。评估或生成样本之前再切到model.eval()。2.2 交叉熵损失错得越多分数越低模型输出的logits不会自动变成“损失”需要一个衡量“错得有多离谱”的指标。对语言模型来说最自然的就是交叉熵损失。PyTorch里用法很简单criterion torch.nn.CrossEntropyLoss() logits model(x) # (B, T, V) loss criterion(logits.view(-1, logits.size(-1)), y.view(-1))为什么要把logits展平因为交叉熵的输入本来是(N, V)每个样本对应一个类别这里每个位置都是独立的预测任务所以把(B, T, V)展成(B*T, V)把标签(B, T)展成(B*T,)即可。有个直观的理解如果词表有500个字符模型纯靠随机猜每个位置猜中的概率是1/500loss大约等于ln(500)≈6.2。所以当你看到loss从6左右开始下降说明训练是正常的如果一开始就远远低于这个数字那就要怀疑是不是数据泄漏了。2.3 反向传播与优化器更新Loss是个标量我们想让模型参数变得更好就得知道每个参数对loss的影响有多大。这一整套求导过程就是反向传播。PyTorch一句loss.backward()就能算完之后每个参数的梯度会挂在param.grad上。接下来要优化器出场。对Transformer类模型我建议直接用AdamW。它把权重衰减和L2正则做了分离Transformer训练几乎默认用它省去很多调参麻烦。optimizer torch.optim.AdamW( model.parameters(), lr3e-4, betas(0.9, 0.95), eps1e-8, weight_decay0.1, )这里几个参数解释一下。学习率3e-4是从零训练小模型的常见起点betas里的0.95是AdamW的momentum比默认更“稳”weight_decay0.1对Transformer来说不算大但能轻微抑制过拟合。不过新手可以暂时不管weight decay的细节先跑通再回来调。反向传播之后必须optimizer.zero_grad()因为梯度是累积在.grad里的。如果不手动清零下一轮backward会把新旧梯度加在一起loss就会变成奇怪的数值。这个bug非常隐蔽我栽过不止一次。2.4 梯度累积与学习率调度实际训练中最常见的痛点是显存不够。一个很有效的应对手段就是梯度累积原来想用batch size32但一次只能跑batch size8那就连续跑4个mini-batch每步都把梯度累加攒够4次再更新一次参数。注意一个细节反向传播时要把loss先除以累积步数再调用loss.backward()。原因是梯度是loss对参数的导数loss放大N倍梯度也会放大N倍。如果你连续四次累加原始梯度最后的更新步等效学习率被放大了4倍训练曲线很可能直接炸掉。除以累积步数之后等效的batch size就是batch_size * accumulation_steps而且梯度的量级和不累积时保持一致超参不需要额外调整。学习率调度我用的是warmup加cosine衰减。训练初期先让学习率从小慢慢升到目标值避免一开始步子迈太大把随机初始化的权重冲散后面再按余弦曲线慢慢降到接近0让loss在训练后期能继续下降而不震荡。total_steps int(len(train_loader) * epochs // grad_accum_steps) warmup_steps int(0.05 * total_steps) def lr_lambda(step): if step warmup_steps: return step / max(1, warmup_steps) progress (step - warmup_steps) / max(1, total_steps - warmup_steps) from math import cos, pi return 0.5 * (1 cos(pi * progress)) scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)注意scheduler.step()要放在optimizer.step()后面而且只有真正更新参数的那一步才调用它。如果每个mini-batch都调用warmup步数会被算多学习率提前衰减训练效果会打折扣。2.5 一个能跑起来的完整训练循环把上面所有东西拼起来一个最简但完整的训练循环长这样def train_one_epoch(model, loader, criterion, optimizer, scheduler, device, grad_accum): model.train() running_loss 0.0 optimizer.zero_grad() for step, (x, y) in enumerate(loader): x x.to(device) y y.to(device) logits model(x) loss criterion(logits.view(-1, logits.size(-1)), y.view(-1)) loss loss / grad_accum loss.backward() if (step 1) % grad_accum 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() optimizer.zero_grad() scheduler.step() running_loss loss.item() * grad_accum if step % 100 0: lr_now scheduler.get_last_lr()[0] print(fstep {step:5d} | loss {loss.item() * grad_accum:.4f} | lr {lr_now:.8f}) return running_loss / max(1, len(loader))有几处容易出问题的地方我特意标注一下。running_loss要做一次* grad_accum因为我们在前面把loss除以了累积步数不乘回来肉眼看到的loss是偏小的。如果训练步数不是grad_accum的整数倍最后几批样本的梯度会一直挂在.grad上不会被更新这会造成每个epoch末尾有数据浪费。最省事的处理是drop_lastTrue把末尾不完整块丢掉或者像我这样不管它影响其实很小。训练循环的本质就是重复“前向、算loss、反向、更新”这四步。你可能觉得它简单但真正让模型学到知识的恰恰是这几行代码背后梯度在数百万参数之间精确流动的过程。3. 让训练更稳的几个实操习惯训练循环能跑不代表能学好。下面这几个习惯是我写各种小模型时逐渐攒下来的每一个都对应我踩过的坑。3.1 固定随机种子代码开头放一个set_seed函数是成本最低但收益最高的工程习惯。import random import numpy as np import torch def set_seed(seed3407): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)为什么重要如果你的数据加载、dropout、参数初始化都不固定每次训练结果完全不同你就会陷入“我改了A发现loss涨了改回去又降了到底谁起作用”的无限困惑中。固定种子并不能保证GPU上的完全复现但至少能让大部分随机源稳定调试起来舒服得多。3.2 梯度裁剪与NaN防护Transformer训练最常见的崩溃方式就是梯度爆炸表现为某个step的loss突然变成nan之后再也救不回来。应付这个问题最直接的手段是梯度裁剪限制梯度的范数上限。total_norm torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)我习惯把返回值total_norm打印出来它比loss更早暴露问题。如果某个step的total_norm达到几百甚至几千说明学习率可能太大、数据有问题或者模型层里数值不稳定。正常训练的total_norm一般波动在个位数或十几以内。如果已经出现nan先别急着调参去检查输入数据里有没有非法值再把torch.autograd.set_detect_anomaly(True)打开定位到第一个爆掉的反向节点。注意这个工具会让训练变慢只用来排查。3.3 日志、验证与Checkpoint训练不是只看loss曲线还需要一个流程化的状态保存机制。否则跑了两天突然断电或程序崩了一切归零那种感觉太崩溃了。我通常每N步打印一次step / loss / lr / grad_norm每个epoch结束跑一次验证集记录valid_loss并保存一个checkpoint。Checkpoint里不只是model.state_dict()还要带上优化器、调度器、epoch、step、best_loss。checkpoint { model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), epoch: epoch, step: step, best_loss: best_valid_loss, } torch.save(checkpoint, checkpoint_last.pt)恢复训练时要把这些state dict全部加载回来。很多人只加载模型参数结果优化器里的动量还停在初始化状态学习率调度器又从0开始训练后面几个epoch的表现会非常怪。3.4 用困惑度和生成样本来判断学习效果Loss本身是个抽象数字我看loss下降看不出“模型到底会不会说话”。所以我一般会额外算一个perplexity困惑度它就是exp(loss)。如果loss是1.5ppl约等于4.5意思是模型平均在5个token里猜下一个token。字符级模型训练充分后能把ppl压到很低因为字符的确定性高子词/词级模型ppl天然会大一些所以不要跨模型比较绝对值只看下降趋势。更直观的办法是让模型实际生成一段文本。生成逻辑很简单输入一段开头模型输出下一个token的概率分布采样一个token拼到输入后面再继续预测。可以加一个temperature来控制随机性def generate(model, tokenizer, prompt, max_new_tokens80, temperature0.8): model.eval() ids encode(prompt, tokenizer.stoi) for _ in range(max_new_tokens): x torch.tensor(ids[-block_size:], dtypetorch.long).unsqueeze(0).to(device) logits model(x) logits logits[0, -1, :] / temperature probs torch.softmax(logits, dim-1) next_id torch.multinomial(probs, 1).item() ids.append(next_id) return decode(ids, tokenizer.itos)训练过程中每隔一段时间生成几条样本看一眼比任何loss曲线都直观。如果生成文本全是重复的词通常是温度太低、模型容量不够或者训练不足。4. 踩坑实录与速查表最后一个部分我把实际操作里踩过的坑整理成一份速查手册。每一条背后都有一段对着屏幕发呆的时光。4.1 Loss纹丝不动现象是loss一直停留在初始值附近比如字符级模型一直卡在6左右不涨也不跌。最常见的原因是学习率太小尤其是当你手贱把lr调成1e-5之后loss就会像被冻住一样。另一个常见原因是数据加载有问题验证一下每个batch的x和y是否是“上一段预测下一段”的关系我曾经把标签错写成了输入本身结果模型直接躺赢loss掉得飞快但完全没学会预测。排查顺序先打印几个batch的x和y肉眼检查是否错位再增大学习率到1e-3跑几十步看loss有没有动静最后看梯度范数如果total_norm一直是0表示反向传播根本没执行到参数上。4.2 Loss瞬间归零或突然NaNLoss瞬间归零别开心往往不是好事。如果刚开始训练loss就从6跳到0.1甚至0说明任务对模型来说太“简单”了简单到它在抄答案。常见原因是数据泄漏比如标签里包含了输入本身的内容或者你验证集划分时和训练集重叠太多。另一种可能是词表太小、语料太少模型几百步内把训练集背下来了。突然NaN的优先级更高它会直接让后续所有step全部废掉。大概率是学习率太大导致梯度爆炸或者数据里有NaN。定位方法在loss.backward()之前检查loss是否有限在optimizer.step()之后检查参数是否有限。assert torch.isfinite(loss).all(), loss is not finite4.3 显存OOM该怎么拆招OOM的第一原则是不要直接缩小模型的block_size然后不管了。缩短上下文确实能省显存但会让模型学不到长距离依赖效果变差。正确做法是把batch size降下来然后用梯度累积把等效batch size补回去。比如原来batch_size32现在改batch_size8、accum_steps4效果基本一样显存直接省一大半。还有几个容易被忽略的显存杀手输入序列没有截断就送去模型、验证时开着梯度计算、跑了太多生成样本没清缓存。验证和生成时记得用torch.no_grad()能省掉大约一半的激活显存。4.4 过拟合、生成复读与早停小模型在小语料上特别容易过拟合表现是训练loss持续下降验证loss却开始反弹。遇到这种情况不要急着把模型做大优先考虑增加数据、增加block_size、增大dropout或者减少训练步数。早停也是一个非常实用的策略每个epoch保存一次当前验证loss如果连续几个epoch验证loss没有下降就恢复此前的最优checkpoint停止训练。生成复读问题大多数时候是采样时temperature设得太低模型被压成了“只挑最高概率的词”或者模型本身还没训练好。把温度调到0.8到1.0之间加上top-k采样通常能缓解复读。4.5 断点续训时状态不同步这一个坑我踩得特别深。训练中断后我重新加载了模型参数继续训练结果loss一下子飙升后来才发现调度器状态没恢复。虽然模型参数一样但优化器里的动量、调度器里的step全都不对等于让一个跑了一半的人突然换成另一个人的腿。正确做法是像前面那样把整个checkpoint完整保存恢复时一次性加载。还要记得给torch.load加map_locationdevice不然在无GPU机器上保存的权重在GPU环境加载时会遇到设备不匹配的问题。4.6 常见问题速查表现象可能原因排查/解决Loss卡住不动学习率过小 / 数据错位打印x与y检查调大lr到1e-3试跑Loss迅速降到接近0标签泄漏 / 任务过于简单检查标签平移扩大语料或词表Loss变成NaN学习率过大 / 输入含非法值 / 梯度爆炸梯度裁剪降低lr检查输入显存OOMbatch太大 / block_size太长 / 未关梯度降batch_size用梯度累积补训练loss降验证loss涨过拟合增加数据或dropout早停回滚生成内容重复温度过低 / 训练不足 / 容量不足调temperature到0.8~1.0加top-k续训loss飙升优化器和调度器状态未加载完整保存/加载checkpoint最后再分享一个我自己的体会。很多人以为训练循环最难的是loss.backward()那一步其实真正决定训练质量的是数据有没有切对、状态有没有管好、日志有没有打全。我一开始也总觉得调库有什么难的直到亲手写过一个完整的训练循环才发现一半时间都花在“确认每个batch的形状对不对”和“恢复断点续训”上。这套迷你GPT虽然离大模型还有很远的距离但数据工程和训练循环的本质是一样的把这条链路彻底跑通后面再往大的方向扩展心里就有底了。
阅读完成 · 觉得有帮助?
咨询建站