做AI工程也有些年了前阵子给自己立了个Flag不依赖任何现成训练框架从零把一个大语言模型的完整流程跑通。项目名就叫 ai-engineering-from-scratch。起因很简单——平时用惯了transformers、DeepSpeed这些库真被问到Attention怎么算的学习率为什么要warmupRLHF到底改的是哪一层参数时发现自己只能答个大概。这让我很不踏实。正好那段时间《Build a Large Language Model (From Scratch)》这本书和build a reasoning model from scratch这类话题也很热大家都在讨论能不能不靠开源权重自己把数据、分词、预训练、微调、推理整条链路走一遍。我决定做一个小规模但五脏俱全的版本用几十GB开源语料训练一个1B左右参数量的基座模型再做SFT和强化学习阶段让模型表现出基本的推理能力。这篇文章就是把整个过程中的设计思路、踩过的坑、以及一些可以复用的经验记下来给也想从头造一遍的同学参考。要做的东西定位很清晰不是要跟Llama、Qwen去比效果而是要理解每一层管线背后发生了什么。所以项目里所有核心模块都是自己实现的HuggingFace只用来做数据读取和评测对照训练部分用纯PyTorch加自定义的分布式逻辑。这套做法适合三类人刚入行想搞懂大模型原理的算法工程师想把业务模型从开源权重切换到自有数据的团队以及被各种框架封装搞晕、想回归本质的开发者。1. 项目到底是什么把从零训练模型这件事拆成可执行的工程1.1 为什么非要从头造轮子先说个现实问题直接用transformersDeepSpeed拉起一个训练任务可能半天就能跑起来那为什么还要自己造轮子我的真实体会是框架封装得太好反而会藏住问题。有一次我遇到loss突然变成NaN定位了半天才发现是某个数据样本里有非法字符如果用的是封装好的Trainer错误信息只会告诉你训练失败你根本不知道它具体在哪一步崩了。自己写训练循环之后每个环节的输入输出、梯度变化、显存占用全都在掌控里出了问题能顺着代码一路查下去。另一方面从零带来的理解深度是看文档替代不了的。当你亲手实现RoPE位置编码、手写FlashAttention的朴素版本、自己调GRPO的advantage计算时你才真正明白为什么这些模块长成现在这样。很多人看模型结构图觉得很简单真到写代码才知道光一个mask的处理就能让你熬夜到凌晨。这也是我坚持这个项目要用手工实现这个方式的原因。1.2 项目范围到什么程度才算从零先说清楚这里的从零不是从晶体管开始。我认可的基础是PyTorch这种自动微分框架、cuBLAS这种底层矩阵库以及开源语料这些属于基础设施没必要重复造。真正要自己动手的是——Tokenizer的BPE合并逻辑、Transformer各层组件注意力、前馈网络、归一化、位置编码、预训练的数据采样器、训练循环含分布式通信、SFT的指令格式、强化学习阶段的奖励建模和策略更新。参考路线上Sebastian Raschka的《Build a Large Language Model (From Scratch)》给了一条很经典的路径先从GPT-2这种结构入手用小型语料验证代码正确性再逐步放大。我的做法和他类似但把重心放在了推理能力这条线上因为这两年推理模型reasoning model很火我想验证一套小规模方案能不能复现让模型在回答问题之前先想一会儿这个机制。这里也多说一句这本书的官方渠道是出版社和作者主页网上那些百度云网盘分享的PDF我不建议用一来版权有问题二来内容可能不完整还是走正规渠道踏实。2. 数据与Tokenizer一切从字节流开始2.1 数据清洗与配比训练效果的隐形决定因素很多人以为数据和模型架构是五五开我实际做下来觉得对于从零训练的小模型数据质量的影响能占到七成。因为模型参数小容量有限喂进去的每一批数据都得是高营养的。我的清洗流程分了四步去重MinHash 精确哈希过滤低质量文本用长度、标点密度、语言模型困惑度做分数过滤去掉个人信息手机号、身份证、邮箱正则匹配最后按领域配比。配比这里有个经验别一股脑把C4、维基百科、代码全倒进去。我给1B模型定的是代码30%、百科与图书20%、网页22%、数学与科学12%、对话数据16%。这样模型既会写代码又有基础世界知识还保留了对话能力。如果代码占比太高模型写代码很溜但日常问答就变得很呆反之代码能力会明显偏弱。这个配比我调了三版才稳定下来每一版都要花掉一次完整的小规模训练来验证。2.2 BPE分词器的手工实现从字符到子词的合并游戏Tokenizer是大模型第一道门也是最容易糊弄的部分。我一开始想着直接用现成的GPT-2词表后来发现要加自定义特殊token比如思维链分隔符时非常别扭干脆自己用BPE算法训了一个32K词表的分词器。核心逻辑很简单先把所有文本转成UTF-8字节序列然后统计相邻字节对的出现频率每次把最高频的字节对合并成一个新token循环执行直到词表大小达到目标。def build_bpe_vocab(corpus, vocab_size32000): # 将语料全部编码为字节序列 byte_tokens [[b for b in text.encode(utf-8)] for text in corpus] merges [] vocab {i: bytes([i]) for i in range(256)} while len(vocab) vocab_size: pair_counts {} for token_list in byte_tokens: for pair in zip(token_list, token_list[1:]): pair_counts[pair] pair_counts.get(pair, 0) 1 if not pair_counts: break best_pair max(pair_counts, keypair_counts.get) new_id len(vocab) vocab[new_id] vocab[best_pair[0]] vocab[best_pair[1]] merges.append((best_pair, new_id)) # 用新id替换序列中所有best_pair byte_tokens [merge_pair(seq, best_pair, new_id) for seq in byte_tokens] return vocab, merges这里有几个坑必须提醒。第一纯Python实现BPE在几十GB语料上根本跑不动我后来改成用多进程分片统计再用字典树做合并速度才提上来。第二分词器训练完一定要在验证集上检查是否有大量未登录词以及平均切分长度是否合理——如果句子被切成太多碎片说明合并轮数不够。第三特殊token在BPE训练时就要预留出来否则后面对不上。我踩过最蠢的坑是训练完才想起来加|im_start|这种token导致词表里凭空多出一堆无意义的字节组合。2.3 采样策略温度、Top-p与数据混入训练时数据加载器不是简单随机采样。我在实践中发现如果完全按原始配比随机抽小模型很容易在某个领域上过拟合而在另一个领域上欠拟合。解决方法是做一个领域加权采样器先按配比给每个领域分配一个概率然后每个batch按概率抽领域领域内再随机抽文档。这样每个batch的领域分布都是稳定的不会因为某个超大语料文件被连续抽到导致训练震荡。另一方面预训练阶段的上下文长度也要渐进式扩展从512开始每训练一段就翻倍到1024、2048。这样做的原因很直接——短序列上收敛快先用短序列把基础语法和知识学好再在长序列上学习长距离依赖能省大量算力。我试过直接从2048开始训同样步数下loss降得明显更慢至少慢了30%。3. 模型架构与训练用代码把Transformer搭起来3.1 架构选型为什么我选了标准Decoder-only再加RoPE模型结构我参考的是GPT-2/Llama这条路线标准的decoder-only TransformerPre-Norm结构SwiGLU激活函数RoPE旋转位置编码Grouped-Query AttentionGQA。选这几个组件不是跟风每个都有明确理由。Pre-Norm加残差连接能让深层网络训练更稳定这是我后面很多问题排查的基石SwiGLU比GELU在相同参数量下效果更好代价是多一组权重矩阵RoPE的好处是外推能力——训练2048长度的模型推理时能勉强外推到4096效果衰减但不会立刻崩溃GQA对推理阶段的KV缓存优化至关重要后面讲推理时你会看到它的价值。注意力部分我建议新手第一版先实现标准softmax attention不要一上来就上FlashAttention。原因很现实FlashAttention的kernel代码复杂如果基础实现还没调通叠加优化只会让bug更难找。我的做法是先实现一个数学校验正确的基础版本跑通后再用PyTorch的scaled_dot_product_attention做替换训练速度直接提升一大截。3.2 训练循环里最容易出问题的几个数值训练循环看起来就是一个简单的for循环前向、算loss、反向、更新。但真正跑起来每个数值都有可能让你怀疑人生。第一个是学习率和warmup。1B模型在bf16混合精度下我的经验是峰值学习率3e-4warmup 2000步然后用余弦退火降到3e-5。关于warmup有个理论解释是训练初期模型的梯度方向噪声很大如果一下子给满学习率容易把参数冲到损失平面的坏区域后面很难回来。这就像开车刚起步就猛踩油门容易撞墙先慢速适应路况再提速才是稳的。第二个是梯度裁剪。我设的max_grad_norm是1.0。这个值不能太大也不能太小太大会失效太小会拖慢收敛。判断方法很简单如果loss在某个step突然跳高多半是梯度爆炸去看梯度的L2范数如果远超1.0就说明裁剪阈值设太大了。第三个是混合精度。bf16训练时loss计算必须用fp32来做否则数值精度不够会导致loss卡住不改。我的loss公式里加了z-loss正则loss ce_loss 1e-4 * logsumexp(logits)这是PaLM论文里的技巧能让softmax输出更平滑训练更稳。加上之后训练的loss曲线肉眼可见地比之前平滑。3.3 让模型学会推理SFT之后的强化学习阶段基座模型训练完成后直接问它数学题通常会得到胡言乱语。这是因为基座模型的目标只是预测下一个词它不是回答问题。所以需要SFT——用人工标注的问题-正确答案-推理过程三元组去微调。这里的关键是数据格式要严谨统一。我用的格式是|im_start|user 一个矩形的长是8厘米宽是5厘米它的面积是多少|im_end| |im_start|assistant |think_start|矩形面积公式是长乘以宽8乘以5等于40单位是平方厘米。|think_end| 所以这个矩形的面积是40平方厘米。|im_end|SFT只能让模型学会模仿推理格式要让模型真正会探索推理路径还得靠强化学习。我用的算法是GRPOGroup Relative Policy Optimization而不是PPO原因很简单——GRPO不需要单独的价值模型省掉一半显存和训练复杂度。它的核心思想是对同一个问题采样多个回答然后根据回答的奖励相对值更新策略。奖励我做了两条答案正确的1格式符合要求有完整思考过程且回答简短的0.5同时用一个小惩罚防止思考部分无限长。跑RL的时候有个现象很有意思初期模型会疯狂输出超长思考因为它发现想得久不一定对但格式分能拿到。这时候要动态调整格式奖励的权重以及限制最大生成长度。我踩过一次大坑思考长度上限设了1024结果模型学会了写循环式思考同一个思路反复绕圈训练速度掉了三倍。解决办法是加一个长度惩罚项超过512 token的思考部分直接扣分。4. 推理优化与部署把训练成果变成可用的服务4.1 KV Cache与批处理推理性能的两个大头模型训练完接下来就是让它真正跑起来。直接逐token生成当然能跑但效率低得可怕。这里必须上KV Cache——把已经算好的Key和Value缓存下来避免每生成一个token就重新计算全部历史。这个优化能把推理速度提升一个数量级。但KV Cache也有代价显存。1B模型、2048长度上下文、GQA设8个KV头每batch的KV Cache大约占用几十MB看起来不多但并发一高就吃紧。这就是为什么架构选型时我坚持用GQA——它把KV头的数量从32个压到8个KV缓存直接缩小到原来的四分之一。推理框架上我自己的实现用来理解原理没问题生产环境还是推荐用vLLM或SGLang它们在连续批处理continuous batching和PageAttention上做得非常成熟同样的GPU吞吐能差3到5倍。4.2 量化与显存控制小卡跑模型的野路子不是所有人都有A100集群。我用两张4090把1B模型跑起来了但部署到一台只有24G显存的机器上时还是需要量化。我的做法是先用INT8量化试水把线性层的权重从fp16转成int8只在反量化时恢复精度。这一步显存减少50%精度损失很小基本看不出来。如果想要更进一步INT4量化GPTQ或AWQ能把模型压到只剩几百MB但我在推理任务上实测数学类准确率掉了2到3个点所以要不要用取决于业务对精度的容忍度。还有个小技巧推理阶段把torch.inference_mode()用起来别用no_grad()。两者功能类似但inference_mode在内存管理上更激进能省不少显存。另外batch size不要拍脑袋定建议用torch.cuda.max_memory_allocated打点实测数据说话别凭感觉。4.3 服务化部署的最小方案从脚本到API部署我选择的是FastAPI加vLLM的组合。vLLM负责模型加载、KV Cache管理和并发调度FastAPI只做一层薄薄的HTTP封装。接口设计上有几个细节值得注意请求要支持max_tokens、temperature、top_p这些参数因为不同业务对生成的随机性要求差异很大要把stream流式输出做进去否则用户等一个完整回答可能要十几秒体验很差还要在服务层做超时控制和错误兜底模型推理偶尔会卡住不能让一个坏请求拖死整个服务。压测下来我这套方案在单张4090上能做到差不多2000 tokens/s的生成速度batch size 16换算成用户感知大约是每秒吐40到80个字基本够用。如果并发再高就得考虑多卡张量并行或者多副本负载均衡了。5. 评估与迭代怎么知道模型真的变聪明了5.1 从Perplexity到准确率评估的层次感训练过程中我最常看的是验证集PerplexityPPL它反映模型对下一个词的预测能力。但PPL下降不代表模型变聪明它只能说明模型学得更顺畅了。真要评估推理能力必须上任务级指标。我给项目搭了一套三层评估体系。第一层是常规基准MMLU知识广度、GSM8K数学推理、HumanEval代码生成、CMMLU中文知识。第二层是格式与结构评估思考过程是否完整、最终答案是否简洁、有没有出现重复循环。第三层是针对业务场景自建的测试集这部分最关键因为通用基准再高业务场景不对路就是白搭。跑GSM8K的时候有个心得小模型对题目措辞极其敏感。同样的数学题数字稍微变一下模型就可能用完全错误的思路求解。所以评估时不能只跑官方测试集得自己做数据扰动——换数字、换人名、换单位这样才能检验模型是不是真的理解了推理结构而不是背题。5.2 用思维链蒸馏和自我反思持续迭代训练完第一版RL模型后我发现它在简单算术上表现不错但一遇到稍微绕弯的题就乱。查了生成日志才发现模型在思考部分经常会写出等等让我重新检查一下这种自我纠正的话这说明它已经在模仿人类的反思行为——这是好事但它的纠错能力不够会说出错误结论后还自信地重复一遍错误。对此我做了一个简单的数据飞轮每评估一轮就把GSM8K上答错的题收集起来人工修正其推理过程混合进SFT数据重新微调再做一轮RL。这个循环跑了三轮GSM8K准确率从52%涨到68%。数据量其实不大每轮只加了2000条左右但针对性极强。这让我深刻体会到小模型拼的不是算力而是数据迭代的效率。6. 常见问题与避坑实录6.1 训练不收敛的典型症状我把这几个月踩过的坑整理成一个速查表希望能帮你少走弯路症状可能原因排查方法loss是NaN数据有非法字符或学习率过高加正则清洗调低lr检查梯度范数loss恒定为某个常数如4.59Tokenizer未收敛模型在输出均匀分布检查词表大小与语料匹配度看logits分布训练正常但生成全是重复词位置编码没生效或温度过低检查RoPE代码调高temperature到0.9多个GPU显存不均衡数据并行分片不均匀用torch.distributed的all_reduce核对batch分配推理时显存暴涨KV Cache未释放或max_tokens设太大监控KV缓存大小改用paged attention这里挑两个最典型的展开说。第一个是loss卡在固定值我排查了很久最后发现是BPE词表训练时特殊token的id没有对齐导致模型每次预测都偏到一个固定位置loss完全学不下去。解决方案也简单重新训练词表确保特殊token在词表训练阶段就预留。第二个是分布式训练时的随机种子问题。数据并行时如果每个进程的shuffle种子不一样不同GPU会吃完全不同的数据导致梯度更新不一致训练震荡。这属于看着没bug实际全是bug的典型。我在数据加载器里用DistributedSampler统一种子后才解决。6.2 我最后悔没早点做的三件事第一件是关于日志。我前期训练监控做得太简陋后来加了wandb之后才发现很多问题其实在训练早期就有征兆loss波动、梯度异常、学习率曲线不平滑。建议从一开始就把loss、lr、梯度范数、显存占用、token吞吐全部打点宁可日志多到看不过来也不要事后追悔。第二件是关于基线对照。刚开始训练时我没有保留一个小型参考模型做同条件对照结果每次改动都不知道是变好还是变坏。后来我固定了一个基线配置特定数据子集、特定学习率、固定步数任何改动都在基线上跑小规模实验做对比效率提升非常明显。第三件是想早点引入自动评估。我早期是肉眼翻生成结果翻了上百条之后发现效率极低。后来写了个简单的评估脚本把GSM8K的每道题、模型答案、判分结果做成HTML报告一眼就能看出哪些题型容易错。这个工具虽然粗糙但让我的迭代速度至少快了一倍。6.3 从项目到经验的总结这个项目我断断续续做了快三个月最终拿到的不是一个多厉害的模型而是一套完整的、可复现的工程方法和无数个原来如此的瞬间。我现在再看各种新出的模型结构、训练技巧不会再觉得它们是黑魔法——因为我知道这些模块在代码层面是怎么工作的也知道它们各自解决什么问题。最后再分享一个小技巧如果你也想做类似的事别贪大先定一个极小的目标——比如用1G语料训练一个100M参数的模型跑通全流程再逐步放大。第一次把全流程跑通带来的信心和手感比任何文档都管用。
