ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

训练MiniGPT实战:从数据加载到文本生成的全流程详解

训练MiniGPT实战:从数据加载到文本生成的全流程详解 训练一个微型GPT模型听起来很唬人但如果你只是想搞清楚大模型从数据到推理的全链路MiniGPT是最好的练手项目。我最近把一套完整的训练流程跑通了从Hugging Face的Dataset加载数据到Context Window怎么切、AdamW参数怎么设再到Training Loop、Validation、Checkpoint和最后的Text Generation前后踩了不少坑。这篇文章把我整理好的方案、代码和问题排查过程全放出来希望对打算入门LLM训练、又不想一上来就怼大模型的同学有实际帮助。1. 整体设计与数据准备先把训练的地基打牢训练一个模型前最重要的事情不是写网络结构而是把数据管好。数据是模型的饭碗饭碗端不稳后面全是白干。这一节我先把MiniGPT的定位说清楚然后重点讲Dataset加载和Context Window切分因为这两个地方最容易出现“看起来没问题一跑就报错”的尴尬。1.1 为什么叫MiniGPT模型结构和参数量怎么定所谓MiniGPT不是某个官方模型而是我们自己定义的一个极小型GPT结构。它保留了GPT的核心设计Decoder-Only、Causal Attention、Token Embedding Position Embedding LayerNorm Feed Forward。常见做法是直接用transformers库里的GPT2Config配置一个小规模的GPT-2比如from transformers import GPT2Config, GPT2LMHeadModel config GPT2Config( vocab_size50257, n_positions256, # 最大pos长度也就是context window上限 n_embd256, n_layer4, n_head8, bos_token_id50256, eos_token_id50256, ) model GPT2LMHeadModel(config) print(f参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M)这套配置跑下来参数量大概在十几M左右一张普通显卡甚至CPU都能扛住。我选择这个规模的原因很简单目标不是刷分而是完整跑通“数据-训练-验证-保存-生成”的整套流程。如果一上来就用几百B的大模型光是显存和训练时间就会劝退大多数人。MiniGPT因为参数少迭代快特别适合做实验也方便我们观察不同超参数带来的影响。1.2 Dataset加载不要踩“closed dataset”的坑数据加载我用的是Hugging Face的datasets库。这个库最大的好处是支持内存映射加载大文件不会把内存吃满而且可以非常方便地做shuffle、split、map操作。我的训练数据是一批短视频评论格式类似于开源社区常见的“douyin comment dataset”那种JSON文件每一行长这样{text: 这个视频太有意思了我已经看了三遍} {text: 求背景音乐名字找了半天没找到}加载代码很简单import json from datasets import Dataset # 如果数据是jsonl格式可以用load_dataset如果是list可以直接from_list with open(datasets/comments.jsonl, r, encodingutf-8) as f: lines [json.loads(line) for line in f] dataset Dataset.from_list(lines) print(dataset) print(dataset[0])这里必须提醒一句datasets库的Dataset对象在使用完后如果手动调用了dataset.close()之后再对数据集做任何操作都会抛出下面这个错误Cannot perform this operation on a closed dataset我第一跑的时候就在验证阶段踩了这个坑。当时我在训练完想释放内存就在dataloader外面顺手加了一句dataset.close()结果后面的验证代码一访问数据就报错。排查了半天才发现是这个原因。所以只要数据后续还要用就不要手动close。如果实在担心内存建议用dataset.save_to_disk()把处理好的结果缓存在磁盘上下次直接加载而不是人为关闭。此外如果你用load_dataset(json, data_files...)这个函数返回的DatasetDict里也包含一个Dataset对象不要提前对它做close。老老实实train_test_split之后让它在内存里待着就行。1.3 Context Window如何切分和构造训练样本GPT训练和普通分类任务不一样它需要把文本转换成固定长度的token序列。Context Window就是模型一次能看到的token数量。我这里是256也就是说每个训练样本是256个token组成的窗口超出窗口的内容要么丢弃、要么通过滑窗截成多个样本。为什么Context Window不能随便设因为Transformer的Self-Attention计算复杂度是窗口长度的平方。256的窗口可能只要几GFLOPs如果拉到1024计算量就会变成16倍。对MiniGPT来说256是速度和效果都比较折中的选择。你也可以用128先验证流程再上256。文本转token后还需要将这些长度不一的token流切成等长的块。transformers官方提供了一种“文本拼接后按窗口切块”的方式核心代码如下from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(gpt2) tokenizer.pad_token tokenizer.eos_token def tokenize_function(examples): return tokenizer(examples[text], truncationTrue, max_length512) tokenized_datasets dataset.map(tokenize_function, batchedTrue, remove_columns[text]) block_size 256 def group_texts(examples): # 将所有token拼接成一个大长串 concatenated {k: sum(examples[k], []) for k in examples.keys()} total_length len(concatenated[list(examples.keys())[0]]) # 计算出可以切成多少个完整block多的尾巴丢掉 total_length (total_length // block_size) * block_size result { k: [t[i : i block_size] for i in range(0, total_length, block_size)] for k, t in concatenated.items() } result[labels] result[input_ids].copy() return result lm_dataset tokenized_datasets.map(group_texts, batchedTrue)有几个细节值得注意。第一labels必须等于input_ids因为GPT是自回归语言模型要在每个位置预测下一个token。transformers的GPT2LMHeadModel会自动完成标签的shift操作所以我们这里把labels直接复制一份即可。第二拼接切块本质上是把不同文本硬塞进同一个窗口窗口内有可能会横跨两个无关句子。但对于语言模型预训练来说这种“噪声”问题不大甚至可以帮助模型学习跨文本的连贯性。第三如果数据很少不建议用太大的block_size否则一个epoch的step数太少模型学不出来。2. AdamW与训练循环模型真正开始学的地方数据准备完毕下面进入核心的训练阶段。这里我重点讲两个东西为什么用AdamW而不是普通Adam以及一个稳健的Training Loop该怎么写。很多新手在网上复制一段训练代码就跑根本不理解每一步在干嘛遇到loss异常只能干瞪眼。2.1 AdamW关键参数不是复制默认值就完事了AdamW本质上是在Adam的基础上把权重衰减和L2正则解耦。普通Adam在做自适应学习率的同时会对所有参数施加L2正则但Adam本身每个参数的学习率不一样这会导致权重衰减的幅度被学习率缩放从而失去正则效果。AdamW则是把权重衰减单独拎出来直接对原始权重做衰减相当于“每个参数都去耦”效果更稳定泛化性也更好。MiniGPT训练时我用的优化器配置如下from transformers import AdamW optimizer AdamW( model.parameters(), lr5e-4, # 对于小模型学习率可以稍微大一点 betas(0.9, 0.999), eps1e-8, weight_decay0.01, )这些参数怎么理解简单说lr全局学习率。过大会导致loss震荡甚至发散过小则训练半天loss下不去。5e-4适合这种百万级参数的模型如果是更大模型建议降到1e-4或更低。betas一阶和二阶动量因子Adam的默认值通常是(0.9, 0.999)一般不用改。eps防止除零的小常数默认1e-8就够。weight_decay正则强度0.01是transformers和GPT-2论文里常用的值。不要拍脑袋设成0.1可能会导致欠拟合。另外我会配合一个warmupdecay的学习率调度器。前一些step把学习率从0慢慢升到目标值让模型在初期不要因为过大的梯度走偏之后再用linear或者cosine衰减。这里用transformers带的get_linear_schedule_with_warmup最方便from transformers import get_linear_schedule_with_warmup total_steps len(train_dataloader) * num_epochs warmup_steps int(0.05 * total_steps) scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepswarmup_steps, num_training_stepstotal_steps, )2.2 Training Loop一个可靠的基础模板训练循环我习惯写成函数这样跑完一个epoch可以单独返回loss也方便后续扩展梯度累积。核心模板如下from tqdm.auto import tqdm import torch def train_one_epoch(model, dataloader, optimizer, scheduler, device, max_grad_norm1.0): model.train() total_loss 0.0 progress_bar tqdm(dataloader, descTraining, leaveFalse) for step, batch in enumerate(progress_bar): batch {k: v.to(device) for k, v in batch.items()} outputs model(**batch) loss outputs.loss loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) optimizer.step() scheduler.step() optimizer.zero_grad() total_loss loss.item() progress_bar.set_postfix(lossloss.item()) return total_loss / len(dataloader)这里每步几个操作的位置不能乱model.train()开启训练模式确保dropout和LayerNorm行为和训练一致。loss.backward()计算梯度。注意这是累积式的所以每次step前必须zero_grad否则梯度会叠加上一轮。clip_grad_norm_梯度裁剪。MiniGPT虽然小但偶尔会出现某个batch的梯度特别大不裁剪容易导致loss突然变成nan。设成1.0是比较稳的。optimizer.step()更新参数。scheduler.step()更新学习率。这一步必须放在optimizer.step之后因为调度器是基于当前优化器状态更新的。optimizer.zero_grad()清空梯度准备下一轮。如果你的显卡显存不够可以加一个梯度累积。比如设置accumulation_steps4每4个小batch才做一次参数更新等效于用4倍batch size训练但峰值显存不变。2.3 训练时的监控和评估准备训练过程不是闷头跑最好每几步打印一次loss和学习率。我一般用tqdm的set_postfix实时显示同时每隔一定step把日志写到文件import logging logging.basicConfig(filenametrain.log, levellogging.INFO) for epoch in range(num_epochs): avg_train_loss train_one_epoch(...) logging.info(fepoch {epoch} train_loss {avg_train_loss:.4f} lr {scheduler.get_last_lr()[0]:.6f})有日志的好处是即便训练中断你也能复盘loss曲线判断是过拟合、欠拟合还是超参有问题。scheduler.get_last_lr()在transformers新版里推荐使用旧版本用optimizer.param_groups[0][lr]也可以。3. Validation、Checkpoint与断点续训很多入门教程只教训练不教验证和保存。但在真实项目里验证和Checkpoint才是决定模型能不能用的关键。没有验证你根本不知道模型是不是在背诵数据没有Checkpoint跑了几个小时突然断电一切归零。3.1 验证集怎么划分评估指标如何计算在数据准备阶段就要把验证集分出来。我用的是train_test_split注意一定要先划分再map否则map后的数据集结构不好处理split_datasets lm_dataset.train_test_split(test_size0.05, seed42) train_dataset split_datasets[train] valid_dataset split_datasets[test] from torch.utils.data import DataLoader train_dataloader DataLoader(train_dataset, batch_size8, shuffleTrue) valid_dataloader DataLoader(valid_dataset, batch_size8, shuffleFalse)验证集的作用是看模型对没见过的数据表现如何。语言模型最常用的指标是Perplexity困惑度它等于loss的指数。一个随机初始化的模型perplexity可能高达几万训练好后能降到几十甚至十几。Perplexity越低说明模型对文本的预测越自信。验证循环和训练循环很像唯一区别是不要更新参数且必须关闭梯度计算def evaluate(model, dataloader, device): model.eval() total_loss 0.0 with torch.no_grad(): for batch in dataloader: batch {k: v.to(device) for k, v in batch.items()} outputs model(**batch) total_loss outputs.loss.item() avg_loss total_loss / len(dataloader) perplexity torch.exp(torch.tensor(avg_loss)) return avg_loss, perplexity.item()两个容易忽略的地方一是验证循环里一定要写model.eval()否则dropout还在工作结果和训练时一样带随机性二是要加上torch.no_grad()否则验证也会构建计算图白白消耗显存和时间。3.2 Checkpoint策略防止训练白跑训练过程中我至少保存两个检查点一是“最新状态”用于断点续训二是“验证集最优状态”用于之后推理部署。最优状态往往不是最后一个epoch因为最后一个epoch可能在验证集上已经开始过拟合。保存函数完整版import os def save_checkpoint(model, tokenizer, optimizer, scheduler, epoch, global_step, eval_loss, save_dir): os.makedirs(save_dir, exist_okTrue) model.save_pretrained(save_dir) tokenizer.save_pretrained(save_dir) torch.save( { epoch: epoch, global_step: global_step, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), eval_loss: eval_loss, }, os.path.join(save_dir, training_state.bin), ) print(fCheckpoint saved to {save_dir})注意这里我同时保存了两类东西model.save_pretrained会保存模型权重和config方便用from_pretrained直接加载training_state.bin保存的是训练过程中的优化器、调度器、epoch等状态断点续训时要把这些全部恢复。加载断点的函数def load_checkpoint(model, optimizer, scheduler, checkpoint_dir): state torch.load(os.path.join(checkpoint_dir, training_state.bin)) model.load_state_dict(state[model_state_dict]) optimizer.load_state_dict(state[optimizer_state_dict]) scheduler.load_state_dict(state[scheduler_state_dict]) start_epoch state[epoch] global_step state[global_step] best_eval_loss state[eval_loss] return start_epoch, global_step, best_eval_loss一个容易被忽略的细节save_pretrained保存的是“模型结构权重”你恢复训练时需要先加载一个相同结构的model GPT2LMHeadModel.from_pretrained(checkpoint_dir)然后再把training_state.bin里的model_state_dict覆盖进去。如果不覆盖你可能会从随机初始化状态开始训练白费前面的时间。3.3 验证阶段容易忽略的坑验证阶段最大的坑不是loss算不对而是验证集和训练集“混了”。如果你在加载数据时忘了做shuffle或者valid_dataset实际上是从train_dataset切片出来的那么验证集里会有训练样本Perplexity看上去很低但真实推理效果一塌糊涂。解决办法是划分时指定seed并固定确保每次运行都一样。另一个坑是数据集的select和map操作可能让Dataset处于某种“惰性”状态如果你在划分后调用了一些会关闭底层arrow文件的API比如dataset.cleanup_cache_files()随后去遍历dataloader就可能出现前文说的Cannot perform this operation on a closed dataset。我的经验是不要在任何dataloader构建之后再对原始dataset做close/cleanup操作。如果内存不够应该先save_to_disk保存然后del原始变量再重新加载。4. Text Generation训练完之后让它说人话训练和验证都通过了接下来就是最有意思的部分让模型生成文本。MiniGPT虽然说不出多惊艳的人话但训练一个下午后它至少能学会评论区的常见句式和用词习惯。这一节我把生成策略和代码完整贴出来。4.1 从Checkpoint加载模型推理阶段不需要再加载optimizer和scheduler直接用save_pretrained保存的目录即可from transformers import GPT2LMHeadModel, AutoTokenizer model GPT2LMHeadModel.from_pretrained(./checkpoint-best) tokenizer AutoTokenizer.from_pretrained(./checkpoint-best) model.eval() if tokenizer.pad_token is None: tokenizer.pad_token tokenizer.eos_token记得model.eval()。同时因为GPT2本身的tokenizer没有pad_token必须在加载后手动设置成eos_token否则后面padding会报错。4.2 生成策略的选择和代码model.generate是transformers封装好的生成函数内部实现了贪婪搜索、束搜索、随机采样等策略。我常用的一组参数input_text 这个视频 inputs tokenizer(input_text, return_tensorspt) with torch.no_grad(): outputs model.generate( inputs.input_ids, max_length64, do_sampleTrue, top_k50, top_p0.95, temperature0.7, num_return_sequences3, pad_token_idtokenizer.eos_token_id, ) for i, output in enumerate(outputs): generated tokenizer.decode(output, skip_special_tokensTrue) print(f生成{i 1}: {generated})这里几个参数的含义do_sampleTrue开启随机采样否则就是贪心搜索每次只选概率最高的词生成结果会比较重复。temperature0.7温度系数。越低生成越保守越高越发散。0.7在稳定性和多样性之间比较平衡。top_k50只保留概率最高的50个候选词再归一化采样避免从长尾里选出太生僻的词。top_p0.95按累积概率截断取累积概率达到95%的最小候选集。和top_k是互补机制两个都开效果通常更好。如果你想要确定性生成可以把do_sampleFalse并设置num_beams5进行束搜索。但束搜索在对话和开放生成场景容易被诟病过于重复我个人更推荐top-k top-p采样。4.3 把生成封装成简单的接口为了后续调用方便我一般把生成封装成一个函数def generate_text(prompt, max_new_tokens32, temperature0.7): inputs tokenizer(prompt, return_tensorspt, truncationTrue, max_length128) with torch.no_grad(): outputs model.generate( inputs.input_ids, max_new_tokensmax_new_tokens, do_sampleTrue, top_k50, top_p0.95, temperaturetemperature, pad_token_idtokenizer.eos_token_id, ) return tokenizer.decode(outputs[0], skip_special_tokensTrue)注意这里用了max_new_tokens而不是max_length。max_length控制的是输入输出的总长度而max_new_tokens只控制新增token数。如果prompt很长用max_length可能会限制生成的长度。还有一个避坑经验训练时的Context Window是256那推理时最好也把输入token长度限制在256以内。虽然GPT的position embedding决定了最大支持长度就是n_positions超过会报错但即使没超输入太长也会占用生成空间。所以我在函数里加了max_length128的截断给生成留出空间。5. 常见问题排查与实操心得最后这部分是实战中沉淀出来的经验我按问题分类整理方便大家直接查。5.1 Dataset操作中的“closed dataset”报错报错信息长这样Cannot perform this operation on a closed dataset常见场景有两个。第一个是前面说的手动调用了dataset.close()。第二个是你在map过程中用了错误的缓存清理方式比如在Jupyter里反复执行可能触发datasets的缓存锁定。解决办法很简单不要手动close如果要在不同数据集之间切换使用copy()或者直接重新加载。如果你确实想释放内存可以先dataset.save_to_disk(cache_dir)然后dataset None等需要时再Dataset.load_from_disk(cache_dir)。5.2 训练Loss不下降如果训练了好几轮loss一直在初始值附近震荡先按顺序检查四件事数据有没有shuffle。如果数据按类别排在一起模型会被连续相似的样本带偏。label是否设置正确。对于GPT2LMHeadModel要保证labels和input_ids尺寸一致并且最好直接复制input_ids。学习率是不是太低或太高。太低学不动太高一上来就震荡。可以打印前几个batch的loss.item()判断。模型有没有在model.train()模式下。如果在eval模式训练dropout被关闭训练动态会不一样。另外如果loss直接变成nan大概率是学习率过大或数据里有脏值。可以先把lr降低10倍再看看是否有空字符串被tokenizer转成了空序列。5.3 显存溢出怎么办MiniGPT虽然小但如果你把Context Window调到1024、batch调到32照样可能Out of Memory。常用降显存手段按优先级排列减小block_size比如从256降到128。减小batch_size然后用梯度累积弥补。开启torch.utils.checkpoint的梯度检查点用时间换空间。使用model.gradient_checkpointing_enable()GPT2LMHeadModel支持。如果还不行就把数据加载放到CPU只把模型和小batch放到GPU或者用混合精度训练torch.cuda.amp.autocast。我用最低配置的显卡4G显存跑256窗口、batch4是没问题的。显存和训练时间是一个跷跷板不要盲目堆batch。5.4 从MiniGPT到真实大模型如何扩展跑通这套MiniGPT流程后转向更大模型其实没有想象中那么难。你只需要把GPT2Config的n_layer、n_embd、n_head调大再把数据规模扩大剩下的训练循环、验证、Checkpoint逻辑基本不用改。如果想做领域微调比如让模型学会某种风格的评论直接用领域数据继续训练即可。进阶方向包括加入定制的tokenizer、引入Instruction Tuning数据做SFT、用LoRA在低显存下微调更大模型以及把生成部分封装成HTTP服务。这些方向的核心骨架都和本文这套MiniGPT训练流程一致。回到开头那句话MiniGPT最大的价值不在于模型本身多强而在于它把大模型训练中最关键的环节全部浓缩在一个下午就能跑完的项目里。Dataset、Context Window、AdamW、Training Loop、Validation、Checkpoint、Text Generation每一个术语都不再是PPT上的名词而是你能亲手改变和调试的东西。最后再分享一个我的个人习惯每次调整超参数我都会把配置写进一个yaml或者dict里并且检查点目录用epoch-02-loss-4.32这种命名方式。这样训练结束后你翻文件夹就能知道哪次实验效果最好。别偷懒Checkpoint命名清晰一点后面能帮你省下大把返工的时间。
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进