☰
从零构建大语言模型:AI工程全链路实战与踩坑复盘
2026/10/3 11:35:13 网站建设 项目流程

1. 先搞清楚“从零开始”到底指什么:范围界定比写代码重要

很多人看到“ ai-engineering-from-scratch ”这类标题,第一反应是:从一片空白里自己写出 tokenizer、写出 Transformer、写出训练框架,最后训出一个能用的大语言模型。这个想象很热血,但也正是这个想象让 90% 的“从零开始”项目在两个月后变成一堆没人再打开的代码仓库。我这次做这个项目时,最先做的不是写代码,而是把“从零开始”这四个字拆开,问清楚自己到底要从哪里开始、到哪里为止。

1.1 工程里的“从零”从来不是“从宇宙大爆炸开始”

如果按照最严格的定义,连 PyTorch 都算不能用的“轮子”,那我们应该从写 CUDA 内核、甚至从设计芯片开始。这当然没必要。工程里的“从零开始”通常有两层意思:一层是“不加载任何别人训练好的权重”,另一层是“不依赖现成的模型训练底座”。前者保证了我们能真正看见模型怎么长出来,后者保证了我们能踩到训练、数据、推理里的真实坑。

我给自己定的范围是这样的:

  • 模型结构代码自己写:Embedding、位置编码、注意力、前馈网络、LayerNorm、输出头,全部用torch.nn的底层原语拼出来;
  • 训练循环自己写:不使用 HuggingFace 的Trainer,手动处理 batch、loss、梯度累积、学习率调度、断点保存;
  • 分词器不自己从头实现:我直接用了tokenizers库的 BPE,因为手写高性能分词并不是这趟旅程的核心价值,但我专门花时间读了 BPE 的合并流程,确保自己明白每个参数在干什么;
  • 权重完全从随机初始化开始:不加载任何开源 checkpoints,这一点没有任何妥协。

这套边界听起来有点“半从零”,但正是这种取舍让项目活了下来。如果一上来就追求“什么都是自己写的”,大概率第一周就被 CUDA 内存分配搞崩溃,根本走不到训练那一步。

1.2 先写 README 再写代码:范围文档的救赎

这个项目开始前,我跟着很多网络食谱式的教程走过弯路——它们喜欢直接甩给你一段代码,然后告诉你“跑起来就行”。但真实工程里,最难的不是跑起来,而是知道自己要解决的问题有多小、有多具体。

我在仓库里先写了一份SCOPE.md,明确写了一句话:目标是训练一个参数量约 1 亿左右的 decoder-only 小模型,在 1 亿 token 规模的中文语料上预训练,之后用少量指令数据微调,让它能在一个垂直场景里生成可读的短文本。这句话直接把很多幻想掐灭了。

为什么刻意把目标设小?因为“large language model from scratch”这个关键词太容易误导新手。你真的复现一篇训练 70B 模型的技术报告,代码可能都一样,但分布式训练、数据并行、模型并行、数十张 GPU 的调度,每一个拿出来都是一个独立课题。我用单卡 24GB 显存,把注意力放在“单机能跑的完整链路”上,这才是普通人可以重复的路径。

定完范围之后,我发现自己心里那个“宏大图景”反而消失了,随之而来的是拆解好的具体任务清单:造数据、写模型、写训练、写看板、写评测、写推理。后面每一步都因为有这个范围而变得可执行。

2. 工具链选型:在“纯手工”和“实用工程”之间找平衡

从零实现最诱惑人的地方,就是一切自己掌控。但工程上有句话叫“不要重新发明轮子,除非你想学习怎么造轮子”。如果你追求的是学习 AI 工程全链路,那么把宝贵的研发时间花在调试 npm 版本、写一个没人维护的 BPE 实现上,并不划算。我做选型时把每个关键依赖都列出来,逐项问自己:它能帮我节省多少时间?它会不会掩盖我本该学习的东西?

2.1 模型与训练框架:选 PyTorch,但不选 Trainer

最终选了 PyTorch 2.x,而不是 TensorFlow 或 JAX。一个重要原因是 PyTorch 的动态图调试体验对初学阶段最友好:你在前向传播里打印一个 tensor 的 shape,看到的和代码写出来的一模一样,这对理解每个维度的流动非常重要。JAX 性能可能更好,但它的函数式风格和隐式编译会让新手在排查问题时多一层心智负担。

更关键的是,我坚持不用transformers库里的模型类和Trainer。Trainer封装了太多细节:数据 collate、梯度累积、日志间隔、断点续训、评估周期……这些确实都是好东西,但它们被“藏”起来了。用Trainer训出一个模型,你依然说不清中间发生了什么。所以我只用它的分词器接口和基础数据集工具,模型结构与训练逻辑全部手搓。

手搓训练循环第一次跑通时很痛苦,但收益很大。你被迫理解每个反向传播、每个optimizer.step()背后的状态更新逻辑。这比从Trainer里拉出一个训练好的模型,要值钱得多。

2.2 分词器:用现成的,但把原理吃透

分词这件事,看起来只是把字符串切碎,实际上却会严重影响模型的表现。我自己在动手前写过一个 100 行不到的玩具 BPE,它把“低”“频”“词”拆成子词的过程,让我对“tokenization”有了具象认识。但真正训练时,我用了tokenizers库的BpeTrainer。

我在选型说明里特别记了一条:不要为了“从零”而在分词上硬造轮子。因为实际训练语料动辄上亿字符,分词速度、内存占用、vocab 文件格式这些坑,只有成熟的库才处理得好。我的项目最终用了一个 32K 词表的中文 BPE,训练语料先跑一遍ByteLevelBPETokenizer,把产出保存下来供缓存复用。这一步省掉了大量重复 CPU 工作。

2.3 优化器、混合精度与日志:默认值不等于最优解

训练一个从头开始的模型,优化器我选 AdamW,这基本是现在大模型训练的事实标准。重点在于学习率和权重衰减:第一个版本我抄了一个技术博客的 3e-4 起始学习率,结果 500 步后 loss 还在 7 附近抖动。后来改成 6e-4 加 cosine schedule,发现下降曲线顺滑很多。

混合精度方面,我推荐直接用bfloat16而不是fp16,尤其在 30 系以上 GPU 上。fp16做 loss scaling 需要多一层 GradScaler 的调参,而且稍不留神就会把梯度打成 NaN。我项目的教训是:前期稳定性比峰值吞吐更优先,等模型能稳定收敛了,再考虑用fp16提速度。

日志这块,很多人看不起“记录每步 loss”,但我吃了亏之后才知道,一条记录着step / lr / gpu memory / loss的日志,是事后排查所有训练问题的唯一线索。我选用了 TensorBoard 做标量可视化,因为它零配置、轻量,不需要额外登录什么平台。

3. 把“大模型”拆成可落地的最小骨架:架构与数据流

模型虽然叫“大语言模型”,但真正可训练的最小骨架其实没那么复杂。我用的结构是标准的 decoder-only Transformer:token embedding 加位置信息,叠若干层 self-attention 和 feed-forward,最后接一个输出头。难点不在结构本身,而在每一层之间的维度流动、mask 处理、初始化比例。今天我把骨架里我认为值得抠的细节展开讲。

3.1 数据管线的第一关:不要把所有文本一次性塞进内存

很多第一次写数据加载的人会直接data = pathlib.Path(...).read_text(),几 GB 的语料瞬间把内存打爆。我项目里语料是一个 12GB 左右的 txt 合集,一次性读进来显然不现实,所以写了流式读取方案:

def read_corpus_lines(path: str): with open(path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if len(line) >= 32: yield line

把文本切成 512 token 的序列时,还有个隐藏问题:如果简单按行切,一条长文本会被截断,语义完整性会受损;如果完全不切,batch 里每条样本长度差异巨大,会造成大量 padding 浪费。我的做法是先把清洗后的文本拼接成“文档流”,然后用滑窗切成固定长度。这样既保证了序列长度一致,又不会丢掉跨行的语义。

3.2 自注意力、因果掩码与 RMSNorm:每个 shape 都要心里有数

我模型实现里比较关键的是自注意力层。输入经过 token embedding 后得到形状(batch, seq_len, hidden_size),线性层生成 Q、K、V,然后 reshape 出num_heads维度。这里最容易错的就是unfold与contiguous()的顺序,我之前因为少写一个transpose,导致注意力矩阵的形状颠倒,却没有任何报错,损失函数照样下降,只是数据排列是乱的。

因果掩码用torch.triu构造一个上三角矩阵,保证第 i 个位置只能看到前 i 个位置。这不是可选项,而是自回归模型的定义。另一个我特意用的结构是 RMSNorm——它比 LayerNorm 少算了一个均值,省掉中间统计量,在长序列训练中更稳定且更省显存。

feed-forward 层我没有用最原始的 ReLU 版,而用了 SwiGLU 风格的门控结构。它的好处是给非线性多了个可学习的门,模型表达能力更强。代价是参数稍多一点、实现多两步。对从零训练的小模型来说,收益大于成本。

3.3 loss 的计算必须显式处理 padding

nn.CrossEntropyLoss默认会对整个序列求平均,这有一个大坑:如果我把 padding 位置也放进 loss,模型会开始疯狂学习“预测 padding 本身”,而不是学会预测下一个真实 token。我在训练循环里手动传ignore_index,让 padding 位置的贡献完全归零。

下面是我训练循环里最核心的一段代码:

# logits: (batch, seq_len, vocab_size) # labels : (batch, seq_len),其中 padding 位置设为 -100 loss_fct = nn.CrossEntropyLoss(ignore_index=-100) shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = loss_fct( shift_logits.view(-1, vocab_size), shift_labels.view(-1) )

shift的意图是:让位置 t 的模型输出,去监督位置 t+1 的真实 token。很多教程会漏掉这一步,导致模型学到的其实是“预测自己”,看起来 loss 也在降,但生成出来的东西是原地复读。

4. 训练循环里的真实一天:损失曲线、梯度爆炸和断点续训

从零训练的前几天,我的日常就是看 loss、改参数、再看 loss。这里没有捷径,但有一套成熟的排查链路。我把这三天里遇到的真实问题列出来,基本能覆盖大多数第一次训练的人可能踩的坑。

4.1 第一次 loss 不为 2.0 时的排查链路

一个刚初始化的小模型,在 32K 词表上,理论初始 loss 应该在ln(32000)左右,约等于 10.37。但我的第一次训练 step 0 loss 只有 4.6,说明模型一开始就“偷看了”什么:后来查出来是 tokenizer 的<pad>token 被大量填充,而 loss 里因为ignore_index设置错误,把 padding 也算进去了,相当于模型预测一个固定 token 就能拿到低 loss。

这种问题不看日志的话完全发现不了。我的经验是:每一步都记录loss / token_accuracy / lr / grad_norm,然后只允许自己在看到这几条曲线的时候调整超参。loss 不降,先检查数据是否对齐;grad_norm 爆了,先检查学习率和梯度裁剪;token_accuracy 长期卡 0.2,基本是 masking 或 label shift 出了问题。

4.2 梯度爆炸、NaN 与 bfloat16:稳定收敛比跑得快重要

训练到 epoch 2 的时候,loss 突然从 2.3 跳到了 NaN。当时第一反应是调低学习率,后来发现根本不是。排查链路是这样的:先看 loss 出现 NaN 的前一步 grad_norm——日志显示数值已经到了 1e5,说明是梯度爆炸。但为什么梯度会爆炸?因为一个 batch 里混入了一段异常文本,里面全是重复的符号,模型在这个样本上的预测极度自信,产生了巨大的 logits。

标准解法是梯度裁剪:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

加上之后,训练稳定了很多。这个max_norm=1.0是经验值,太小会让模型收敛变慢,太大的话等于没裁剪。我还做过一个实验:把 optimizer 换成 SGD,发现收敛速度慢得离谱,最终又换回 AdamW。

还有一次 NaN 的根源更隐蔽:一个前向 tensor 经过bmm之后出现 Inf,因为中间值超过 float16 的表示范围。我的建议是:如果显存够用,优先用bfloat16而不是fp16;如果必须用fp16,那autocast和GradScaler都要显式开,别省。

4.3 断点续训不是简单保存模型权重

训到第 7 个小时,云服务器因为维护被重启了一次。我原来的 checkpoints 只保留了模型权重,结果恢复训练时 optimizer 的动量状态全丢了,学习率调度也重新开始,相当于浪费了大半天。之后我改成每次保存五个状态:模型权重、优化器状态、学习率调度器状态、RNG 状态、当前 step 和数据采样位置。

保存 RNG 状态这一点很多人注意不到,但数据加载如果依赖了 shuffle,不保存 RNG 会导致恢复训练后数据顺序不一致,损失曲线会出现一个跳变。我用的简化版保存代码大概是:

checkpoint = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "step": step, "rng": torch.get_rng_state(), } torch.save(checkpoint, f"ckpt-{step}.pt")

验证恢复是否成功的方法也很简单:保存后继续跑 100 步,再重启,看看新 100 步的 loss 是否与之前那条曲线平滑衔接。如果出现“折线式”断崖,说明状态没保存全,返回去检查 list。

5. 推理与落地优化:从“能跑”到“跑得值”

模型训出来了,不等于工程完成了。我花了比训练还多的精力做推理端优化。原因很简单:训练只要看曲线,推理要面对面听模型胡说八道。这也是“ ai-engineering ”里最容易被低估的部分。

5.1 自回归生成与 KV Cache:速度翻倍的秘密

生成文本时,最朴素的做法是每一步把 history 整体重新喂进模型。这在概念上没错,但对同一个已经生成的前缀,Q、K 都被重复算了几十遍。KV Cache 的思路是:第 i 步算出的 K 和 V 缓存下来,第 i+1 步只计算新 token 的 K、V,然后拼接上去。对一个 12 层 Transformer 来说,这个改动能让生成速度从“看着它一个字一个字蹦”变成“接近打字机输出”。

实现 KV Cache 时要注意显存开销:cache 的大小约等于batch * seq_len * hidden_size * num_layers * 2(K 和 V),序列越长越吃显存。我项目里把 cache 显存上限写进日志,防止生成很长文章时 OOM。

我还做了两个工程细节:生成过程中如果碰到<eos>,马上跳出循环;设置了最大生成长度 256 token,防止模型陷入无限循环。这些都是光看教程学不到的,因为教程只会给你一个 20 行的 demo,不会管你真实跑一个长文本生成时会被无限重复文本卡多久。

5.2 采样参数:temperature、top-k、top-p 为什么要一起调

第一次跑通生成接口,我直接用argmax采样,模型输出非常死板,每句话都是高频词组成。后来加了 temperature=0.8,灵活很多,但又开始出现不通顺的词。之后我又加了 top-k=50 和 top-p=0.9,才算平衡了“多样性”和“通顺度”。

这个组合没有绝对最优,需要针对你的模型、语料和用途试出来。我的经验是:小模型更适合低 temperature,比如 0.7 到 0.8,因为它的概率分布本身就比较平,temperature 再高就会碎。大一点的模型可以稍微放宽。每次改完采样参数,别只看一句输出就下结论,我会在固定 prompt 下连续生成 10 条,人工看一遍重复率、语法错误率和语义离题率。

5.3 真正部署时看到的另一层问题

我最后用 FastAPI 写了一个最简单的生成服务,然后立刻踩到一个新手坑:推理时忘了model.eval()和torch.no_grad(),模型的表现忽好忽坏,因为 dropout 还在影响前向传播。这个问题很典型,因为它不报错,只有对比多个请求后才会发现“为什么同样 prompt 两次生成差别这么大”。

部署端的量化优化我也做了一步:把权重从 fp32 转成 bfloat16 再加载,显存占用减半,生成速度小幅提升,质量几乎不受影响。再往下做了torch.compile加速,第一次编译很慢,但之后生成的每 token 延迟少了 20% 左右。这种收益算不上惊天动地,但足以看出工程优化的复利效应。

6. 用自己的模型说一句人话:从第一句话到达标的全流程复盘

最后一个部分,我想给这个项目画个完整的句号。因为这篇文章的灵感来源,其实是一个很网红的标题:“build a reasoning model from scratch”。我对这个标题有很复杂的情绪。它听起来像是一个下午就能完成的事,实际上包含了预训练、指令微调、甚至增强学习的完整链路。作为走过一遍的人,我想聊聊什么才是合理的验收标准,以及“推理能力”到底能不能从零长出来。

6.1 给模型设一个不羞耻的验收标准

1 亿参数的模型,你不可能期待它像 GPT-4 一样写代码、讲段子。我给自己的验收标准是:输入一句“北京是中国的首都”,模型有没有能力把这句话续写完整;输入一个短问题“1 加 1 等于几”,模型能不能给出接近的答案形态。这些标准听起来很低,但对一个完全从头预训练的模型来说,已经是很多不稳定训练的产物了。

我记得第一次让模型输出内容时,它生成了一句“北京是中国的首都,是国家的政治中心”。翻译成 token 人类看起来甚至有点通顺,可能是因为语料里这类句子太多了。那一刻的真实感受比跑通任何 benchmark 都强烈。这就是“ from scratch ”带给你的认知:你不再觉得模型输出是魔法,而是每一步前向传播叠加出来的统计结果。

6.2 推理能力不是代码能“带”出来的,数据才是开关

网上那些“从零构建推理模型”的帖子,大多只是搭了一个模型框架,不会告诉你:在通用文本上预训练得到的模型,你说一句“请一步一步思考”,它并不会真的推理。它只会把“一步一步”“首先”“然后”这些词的高频搭配背出来。真正让模型表现出推理迹象,靠的是训练数据里确实包含大量成体系的思维链样本。

我做了一次对比实验:拿同一个基座模型,用一份 2 万条左右的中文 CoT 样本做微调,再用相同的模板提示它。微调后,模型开始会输出“首先根据已知条件列出变量,然后……”。虽然这个“推理”在逻辑上很脆弱,但格式和顺序确实像那么回事。这个实验让我明白:所谓 reasoning model,工程端要解决的是数据配比、格式设计和训练策略,而不是一上来就再造一个架构。

6.3 账单、时间与最后的一点建议

整个项目最贵的不是机器,而是时间。我用的是单张 24GB 显存的 GPU,预训练阶段跑完整轮约 60 小时,加上调试期间反复重启,总时长超过 100 小时。按市场租赁价格折算,训练成本大概是几百到一千元级别。这个成本对于一个人体验“从零构建”来说,是完全可以接受的。

如果你也想复现,我给三个建议。第一,把小模型、小语料、小目标作为第一次迭代;不要复制小说级规模,否则你会在一个完全没必要的地方烧掉时间和信心。第二,每一行训练代码都必须知道它在干什么;不要因为别人推荐了Trainer就直接用,它会让你以为“从零训练”不过如此。第三,把失败日志当作收获;你后面吹牛时最值得讲的不是模型最终效果,而是你从 NaN、断点、数据错位里爬出来的全过程。

最后分享一个细节:我到现在还保留着第一次训练跑到 loss 降到 3.0 时记录的日志文件。每次看到那条曲线,都能想起整个调试过程的焦灼。这个项目教会我的是,AI 工程里“从零”的真正含义,不是不借助任何现成工具,而是不借助任何模糊的理解。你愿意亲手拆开每一个环节,拆到能对别人讲清楚它的输入输出和踩坑点。能做到这一步,你的“ from scratch ”才真正算数。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询