做AI工程这行很多朋友一上来就接个大模型API跑通demo就觉得自己会了。但真正踩过从零构建的坑之后你会发现“AI engineering from scratch”六个词背后是数据、训练、推理、部署一整条看不见的链路。这篇文章是我近期一次完整复现经历的复盘从空的代码仓库出发构建一个可运行的、约一亿参数的自回归语言模型包括数据管线、tokenizer训练、训练循环、KV cache推理优化和最终服务化部署。它适合两类人——一类是“只调API没碰过权重”的工程同学另一类是打算系统入门大模型底层但被资料淹没的学习者。我会尽量按实操顺序讲把选型理由和踩坑过程一并交代清楚有些数字和代码你可以直接抄走。1. 先想清楚这个“从零”到底从哪开始1.1 为什么我坚持走一遍最难的路现在大模型开发最不缺的就是框架。HuggingFace transformers几行代码就能加载一个几十亿参数的模型vLLM一条命令就能拉起来一个推理服务看起来“从零构建”这种行为确实有点自讨苦吃。但我自己的体会是框架把太多关键细节封装在了黑盒里你只知道调用model.generate()但不清楚KV cache在它内部是怎么更新和释放的你只知道trainer.train()能跑起来但不明白数据加载顺序、loss掩码、优化器状态恢复这些环节是怎么串起来的。一旦到了生产环境模型要接入真实业务要改结构、要调性能这种“只知其然”的状态就会成为瓶颈。我自己就遇到过必须魔改注意力头来支持特殊输入格式的需求那时候翻内置源码的体验和真正自己写过一遍神经网络完全不是一回事。所以我定下一个原则这次只借助最底层的PyTorch和必要的tensor库模型结构、训练流程、采样逻辑、推理缓存全部手写。这不是为了炫技而是为了把每一个设计决策变成自己真正消化过的东西。当然我也建议读者根据目标取舍如果你是做LLM应用、追求快速迭代直接用现成框架完全合理如果你想把底层原理吃透或者需要频繁定制结构那这条from scratch的路值得走一趟。1.2 项目全貌一个1亿参数的decoder-only模型我先交代这次项目的完整边界。目标不是做一个能跟GPT-4掰手腕的模型那超出个人工程项目的合理范围我瞄准的是一个约1亿参数的自回归语言模型训练数据规模控制在10亿token量级最终能完成流畅的中文续写、简单问答和风格模仿。参考《Build a Large Language Model (From Scratch)》这类资料的路线我把它拆成四个阶段。第一阶段是数据工程包括语料清洗、去重、切分、编码以及自己训练一个BPE分词器。第二阶段是模型搭建与训练包含完整的Transformer decoder结构、损失函数、优化器、学习率调度、混合精度训练和断点续训。第三阶段是推理引擎包括自回归生成、采样策略、KV cache缓存、动态量化。第四阶段是服务化部署用一个轻量级HTTP框架把模型包装成真正的线上接口。这四个阶段听起来清清楚楚但实际操作中会有大量环路跑起来发现loss不收敛要回头查数据清洗和label对齐部署时候显存吃紧要回模型里调整结构。所以后面我会按“数据 → 训练 → 推理 → 部署 → 排错”的顺序来写排错这章是我认为最值钱的部分很多坑靠官方文档根本查不到。1.3 给项目定一个可验收的目标没有验收标准的工程很容易烂尾。我的验收指标有三条全部用数字说话。第一验证集loss要低于某个明确阈值。以1亿参数和10亿token量级来说我把目标定为中文验证集困惑度低于40换算成交叉熵大约在3.7以下这个数字不激进但足以证明数据管道和模型实现是自洽的。第二单次生成效果要可主观评估给定“春天到了我推开窗”这样的开头模型至少能作出三句以上语义通顺、不明显重复的续写。第三是性能指标在单张A10或RTX 3090上生成128个token的总延迟不超过1秒服务端能支撑至少2个并发请求不超时。我后面所有优化手段本质上都是围绕这三个目标进行的。提示验收标准一定要在动手之前写下来。否则很容易陷在“再驯一轮loss”的循环里永远发不了版。2. 数据工程模型能不能学好的第一道关卡2.1 语料选型与清洗策略很多初学者把注意力放在模型结构上但实际上数据工程决定了模型的天花板。我这次选择公开的、来源合规的语料。具体以中文维基dump为主混入若干开源中文文章集整体原始文本大约20GB。拿回来之后不能直接喂给模型必须过一套清洗流水线。我的清洗顺序是这样的先把原始dump里的HTML标签、模板信息、页面导航这些非正文内容全部剔除这一步用正则和简单状态机配合完成开源库里有现成工具可以借鉴但建议自己过一遍逻辑因为不同dump的格式差异很大。接着做归一化统一全角半角、压缩连续空白、去掉控制字符这里注意不要把换行符误删否则后续分句会很痛苦。然后做质量过滤规则包括去掉长度小于50字的文档去掉字母数字占比过高或符号占比过高的段落去掉广告、导航、免责声明等模板化文本。最后做去重我用了MinHash算法计算文档指纹近重复内容直接丢弃这一步能省下大量无效训练token。清洗完之后我还会做一个耗时但必要的步骤把全量数据手动抽样看一遍每5000条里抽1条看看有没有规律性问题。实测中我确实发现有一部分文本是乱码拼接这种问题靠统计指标很难发现只有人眼才能识别。2.2 tokenizer自训BPE与词表设计大模型的tokenizer决定了模型看世界的基本单位。我选择使用BPE字节对编码因为它在中文场景下天然优雅中文单字频繁出现可以直接合并成词级token而遇到生僻字、英文字符、数字时又能退到字节级避免出现大量[UNK]。训练BPE时我先做了一版基线实验语料分别用30MB、300MB、3GB训练词表观察相同文本下的分词稳定性。最后我选择用3GB采样语料训练一个大小为32000的BPE词表并添加了bos、eos、pad、unk这四个特殊token。这个大小在中文场景下比较合适词表太小中文会被切得很碎序列过长词表太大embedding矩阵占用的显存会明显上涨以d_model768计算词表每增加10000embedding参数就增加768万训练速度也会跟着吃紧。实现BPE时有一个重要细节先做NFKC归一化再进行预分词否则全角英文和半角会分裂成两种token浪费词表容量。另外训练语料的采样比例要保持多样尽量覆盖不同主题和写作风格否则词表会偏向某一个垂直领域。2.3 数据规模、切分比例与缓存格式数据规模直接影响训练步数和成本。我先算一笔账目标是训练10亿token模型参数约1亿比例大约是10比1。这个比例在学术界不算高通常建议20比1甚至更高但考虑到个人项目的算力限制10比1已经能产出一个行为正常的模型所以我选择在此基础上跑足。有了10亿token之后我把数据按99:1拆成训练集和验证集验证集严格不参与训练。这里有个很容易犯的错误有些人按“文档数量”切分而不是按“token数量”切分会导致验证集和训练集主题分布不一致。我的做法是把所有文档随机打乱后用流式计数器统计token数到总token数的99%处为训练集剩下1%为验证集。这样能保证两边文本分布相似。缓存格式我推荐用mmap文件把预处理后的token序列以uint16数组形式写入磁盘。uint16最多能表示65535个id完全够用。相比parquet和TFRecordmmap在读取时不需要把整个文件load进内存数据loader可以随机访问任意位置这对训练效率很关键。Dataloader部分我简单贴一个核心逻辑class TokenDataset(torch.utils.data.Dataset): def __init__(self, mmap_path, seq_len, n_tokensNone): self.data np.memmap(mmap_path, dtypenp.uint16, moder) self.seq_len seq_len self.n_tokens n_tokens or len(self.data) def __len__(self): # 每个样本为一个长度为seq_len1的窗口最后一个是预测目标 return max(0, (self.n_tokens - 1) // self.seq_len) def __getitem__(self, idx): start idx * self.seq_len chunk self.data[start: start self.seq_len 1] x torch.from_numpy(chunk[:-1].astype(np.int64)) y torch.from_numpy(chunk[1:].astype(np.int64)) return x, y这里x是输入y是右移一位的目标。要注意最后一个样本可能不足seq_len1我会做截断或直接跳过避免参差不齐导致显存碎片。2.4 流式训练中的shuffle策略常规机器学习里shuffle是把所有样本打乱后重新切batch。但大模型训练时数据动辄几十GB甚至上TB全量在内存里做shuffle是不现实的。我的做法是分片shuffle先把整个语料切成若干个固定大小的分片每个分片内部打乱再按分片轮流读取。这样可以保证同一个epoch内不会连续出现同一主题的长文本同时内存消耗可控。另一个细节是epoch边界处理。训练时我按照“全局步数”而不是“epoch步数”来管理进度这样断点续训时只需要记录数据读取器的偏移量和随机数种子即可。如果不处理这个细节中途恢复训练会重复读一部分数据导致模型重复学习某一段语料影响最终loss。3. 模型架构与训练实现3.1 现代decoder-only架构的组件选型这次我手写的模型是一个标准的decoder-only Transformer但不再是经典论文里那种朴素结构而是把现代大模型常用的改进都用上了。核心模块包括token embedding和输出层共享权重、RoPE旋转位置编码、RMSNorm、SwiGLU激活函数、因果多头注意力以及可选的Grouped Query Attention。我逐个说为什么选这些组件。第一RoPE替代了原始Transformer里的绝对位置编码。绝对位置编码把位置信息加在词向量上而RoPE是旋转矩阵作用在Q和K上它天然保留了相对位置关系而且对外推更友好。我实测里RoPE模型在长文本生成时的稳定性明显优于正弦编码。第二RMSNorm比LayerNorm少算均值只做缩放效果反而更稳定。第三SwiGLU是一种门控激活函数计算量比ReLU略大但能带来可感知的收敛收益。第四GQA不是为了提升效果而是为了省显存和推理速度它让多个Q共享一组K和VKV cache体积能缩小到原来的若干分之一。我的具体config参数如下参数值n_layer12d_model768n_head12n_kv_head4 (GQA)d_ff3072vocab_size32000max_seq_len512activationSwiGLUnormRMSNormpos_encodingRoPE这套配置参数量大概在1.05亿左右单卡显存占用可控又足够表现出大模型的基本行为模式。3.2 参数量、算力与显存的预先估算动手训练之前必须先把三件事估算清楚参数量、计算量、显存峰值。参数量公式不复杂embedding矩阵是vocab × d_modelTransformer block里注意力、FFN、norm各有自己的参数。我按上面的config估算参数量约1.05亿其中embedding部分约2457万Transformer主体约8000万。计算量用经典的近似公式FLOPs ≈ 6 × N × D估算N是参数量D是训练token数。代入N1e8D1e9得到约6e17FLOPs也就是0.6 ExaFLOPs。以单张A100的312 TFLOPS理论算力计算如果利用率能做到50%大约需要3800秒但实际训练时很难持续达到50%利用率我跑下来真实耗时大约是理论值的2.5到3倍所以一周左右的时间预算才算合理。显存是最容易把卡顶爆的一环。我把公式拆成四部分模型权重、梯度、优化器状态和激活值。使用BF16混合精度训练时权重和梯度各占2 × N字节AdamW不仅保留额外一份FP32权重副本还维护一阶和二阶动量这部分占8 × N字节激活值取决于batch size和序列长度往往才是最大的不确定项。按batch size 64、序列长度512、总token数32768来估算激活显存接近4GB。我把这些数字整理成下表项目估算显存模型权重BF160.21 GB梯度BF160.21 GBAdamW状态FP321.26 GB激活值batch 64 × seq 5123.5~4.5 GB合计约5.5 GB也就是说单张24GB的3090完全够用。如果batch size再翻倍激活值一起吃紧就要考虑梯度累积或者activation checkpointing。3.3 训练循环、超参与监控训练部分最关键的是把学习率、优化器和损失函数这几个点调对否则模型即使能跑loss也会飘在半空不下降。优化器我选了AdamWlr3e-4、weight_decay0.1、beta(0.9, 0.95)。其中weight_decay只作用于权重矩阵不作用于norm层和embedding这样可以减少过拟合风险。学习率调度用“warmup cosine decay”前500步线性上升之后按余弦曲线逐步衰减到峰值学习率的10%。这个策略的性价比很高又稳又能让loss收敛到较好的极小值。损失函数本身和做分类一样用交叉熵但有一个细节必须注意在batch里有一些padding位置它们的预测目标要设置成ignore_index-100这样loss计算时会自动跳过这些位置否则模型会努力去学习“预测下一个pad”白白浪费学习能力。混合精度我选择了BF16而不是FP16。BF16的数值范围和FP32一样大只是尾数精度低在loss较大时不容易溢出FP16在小数值场景下容易出现NaN。当然BF16在部分老显卡上不支持如果你只有V100或更早的卡建议还是走FP16并配合loss scaling。训练监控是另一个重点。我每个step记录四类信息train loss、validation loss、当前学习率、梯度范数。梯度范数是判断训练健康度的关键指标如果它突然暴涨到几十甚至上百说明某一步出现异常这时候应该停止训练排查而不是硬跑下去。每2000步做一次验证每5000步保存一次checkpoint。checkpoint不只要存模型权重还要把optimizer状态、学习率调度器位置、数据读取器偏移量一起存进去否则断点续训时优化器状态从零开始几乎等于重新训练。4. 推理优化与部署4.1 自回归生成与采样策略模型训练完成后一个完整的推理链路至少包含编码、自回归生成、采样、解码四步。自回归的机制很简单模型输入一整段prompt预测出下一个token的概率分布然后根据采样策略选出token拼接到输入尾部重复这个过程直到生成eos或达到max_new_tokens。但直接“每次选概率最大的token”贪心解码会带来一个问题文本容易陷入重复循环。我尝试过让它续写“今天天气很好”模型会重复输出“阳光洒在窗台上”好几遍这就是贪心解码的经典现象。所以实际接口里我默认使用采样先对logits除以temperature再叠加top-k和top-p过滤。temperature0.7是我试下来比较稳的默认档位既保留一定创造性又不会太散。此外还加了repetition penalty每次生成时对历史出现过但非bos的token施加一定的惩罚系数能明显压低重复率。def sample_next_token(logits, temperature0.7, top_k50, top_p0.9, penalty1.1): logits logits / temperature # repetition penalty logits[history_ids] - penalty # top-k if top_k is not None: k min(top_k, logits.size(-1)) indices logits.topk(k).indices mask torch.full_like(logits, float(-inf)) mask.scatter_(-1, indices, 0) logits mask # top-p if top_p is not None: probs torch.softmax(logits, dim-1) sorted_probs, sorted_indices probs.sort(descendingTrue) cumsum sorted_probs.cumsum(dim-1) remove_mask cumsum - sorted_probs top_p logits[sorted_indices[remove_mask]] float(-inf) return torch.multinomial(torch.softmax(logits, dim-1), 1)这段代码可以直接抄走注意penalty会改变logits概率分布所以要在softmax之前做。4.2 KV cache把重复计算省掉推理优化的第一课就是KV cache。在自回归生成中每生成一个新token模型都要重新计算一遍所有历史token的注意力。如果没有缓存第100个token的生成也要重新算前99个token的K和V这等于不断地重复做矩阵乘法越到后面越慢。KV cache的核心思路是已经计算过的历史token的K和V在后续步骤中其实不会变那么就在显存里把它们存下来每次生成新token只需要为当前token计算一次K和V然后拼接到缓存里再和Q做注意力即可。实现上需要注意两个边界第一是prompt预填充阶段生成的KV可以直接填满缓存第二是缓存必须按batch、n_kv_head、seq_len、head_dim的维度预先分配注意head维度是分配max_new_tokens prompt_len还是max_seq_len要根据业务场景提前预留。我改造模型时把注意力层换成一个带cache的版本核心差异只有几行。效果则非常明显生成128个token未缓存版本耗时约1.8秒加上KV cache后降到0.9秒速度几乎提升一倍。如果模型更大、生成更长这个优势会进一步放大。4.3 量化实践与HTTP服务封装KV cache解决的是“重复计算”量化解决的则是“跑不动”的问题。我这次主要尝试了动态INT8量化把每层的权重矩阵按行列的absmax映射到INT8整数范围推理时先把权重临时反量化回FP16参与计算。这样做的好处是不需要校准集实现简单缺点是内存带宽收益有限因为每层计算前都需要反量化。在真实部署中GPTQ和AWQ这类静态量化效果更好可以对权重做更细粒度的量化并能直接在INT8上做部分矩阵乘法。但实现复杂度高出不少。对于我这个1亿参数规模模型INT8动态量化本身意义不算大权重大小从200MB降到100MB也不能改变推理瓶颈在注意力计算这个事实。所以这里我的经验是小模型优先做剪枝和KV cache优化大模型才开始考虑复杂量化方案。部署层我用FastAPI封装了一个简单的/generate接口接收prompt和采样参数返回生成文本。服务进程加载一次模型权重用一个全局锁控制并发避免多线程同时解码导致显存爆炸。上线后我压了一阵子单机单卡支持3到4个并发请求单请求延迟约0.7秒。如果想进一步提升并发吞吐我会推荐改用vLLM或TensorRT-LLM作为推理后端这些框架已经内置了连续批处理continuous batching和PagedAttention但不属于本次“从零实现”的范围。5. 常见问题与排查实录5.1 训练期loss不降、发散与显存溢出我把它做成一份速查表每一条都是自己踩过之后验证有效的。现象常见原因排查与解决loss完全不降学习率过大或过小数据与label错位embedding未初始化先以1e-4小lr跑百步若loss有下降趋势再加大检查y x[1:]切片是否真对齐用0.02标准差初始化embeddingloss突然变成NaNBF16/FP16溢出学习率峰值过高换BF16并依赖其宽数值范围把warmup步数增加避免一开始就冲高loss下降但验证集飙升训练数据泄漏到验证集过拟合检查去重逻辑是否覆盖验证集增大weight_decay降低训练epoch数显存OOMbatch过大序列过长激活值堆积减小batch配合梯度累积保持等效batch大小不变启用gradient checkpointing降低max_seq_len还有一个经常被忽略的问题如果你的模型和某开源实现行为对不上比如同数据下loss差出一大截千万不要瞎调超参先写一个脚本把每一层的输出形状和梯度范数打印出来逐层对比参考实现通常能快速定位到某个矩阵转置顺序或者mask实现错了。5.2 推理期生成重复、越界与延迟超标推理期的三大问题分别是生成内容重复、KV cache索引越界、单请求延迟太高。生成重复是最常见的。不要一味调大repetition penalty因为惩罚太强会导致模型翻来覆去只换几个词。我建议先调高top-p到0.95再调temperature到0.8最后才考虑penalty。如果仍然重复就去检查训练数据里是否有大量重复模板很可能模型学到的是“复制粘贴”这种消极策略。KV cache越界通常发生在prompt长度加上max_new_tokens超过缓存预设长度时。排查方法很简单把生成函数包进try/except打印缓存数组的shape和访问索引一眼就能看见是不是越界。为了避免这种问题我的经验是缓存统一按max_seq_len分配也就是把seq_len和max_new_tokens共用同一个上限而不是分开算。延迟超标的原因很多但排除下来最常见的是没有用KV cache其次是没有做batch推理。很多人会疑惑“我明明就一个用户请求为什么还要batch”答案在于如果服务端同时接收多个请求单独逐条推理会重复加载权重、重复分配显存。用动态batch把不同长度的请求拼在一起padding不仅能提高吞吐还能摊薄固定开销。5.3 我的几条避坑心得这些是一条条踩出来的经验很难在教科书上直接找到我单独列出来。第一全流程先在“微型语料”上验证再上全量数据。我第一次动手时直接拿20GB语料去训练结果跑了两天发现BPE把\n和普通字符合并出了大量垃圾token只能全部重来。后来所有朋友我都建议先拿20MB语料、跑100步验证整个链路能走通再做全量训练。这个习惯至少帮我省下了一周的算力预算。第二把tokenizer、数据管线、模型结构、采样逻辑分别做成独立模块并写对应的单元测试。具体来说tokenizer的“编码再解码”要保证能还原原文数据管线的(x, y)要保证长度和偏移量逻辑正确模型的loss要和参考实现对比过才算通过。这一步很多人嫌麻烦但正是它让后面所有调试都变得有迹可循。第三每个训练实验都记录config、seed、数据版本、commit号。没有这套记录loss出现波动时你根本无从排查是代码改动还是随机因素导致。我自己吃过一次大亏一个数据清洗规则的改动没有记录导致后续十几个实验和基线对比完全失去意义。第四不要一开始就上分布式训练。单卡把整个流程跑通已经足够学会90%的工程细节分布式框架DeepSpeed ZeRO、FSDP只会额外引入通信和调度的复杂度。等你发现单卡真正成为瓶颈再引入分布式你会更容易理解它到底解决了什么问题。到这里我想多说两句这几年越来越多人把大模型当成一个黑产品来用但“从零构建”这件事给我的回报并不是那个可以续写中文的模型本身——那东西在生产环境里其实不稀缺。真正的收获是我现在敢对模型结构、数据管线、推理策略做任何程度的修改因为我知道每一个模块内部是怎么转的知道改动之后会在哪个环节产生连锁反应。如果你也想复刻这条路线我个人最强烈的建议是先把规模压下来用200M参数和1B token把全流程跑通再往大扩展。下一步我自己准备尝试的是把同样的思路应用到构建一个能进行简单推理的reasoning model上核心不是训练方法本身而是数据构造和验证反馈闭环。希望这篇复盘能帮你少踩几个坑把精力留给真正有意思的部分。