Happy-LLM 学习笔记 06：动手搭建一个小 LLM，最值得盯住哪些模块

这一章最有价值的地方，不是把一个完整的大模型“抄出来”，而是把一个小 LLM 从骨架、入口和训练三层拆开看。以前我总会先盯参数量和显卡，后来才发现，真正决定能不能跑通的，往往是模块之间的接口是否统一。

先把模型骨架搭对

k_model.py 里最先出现的是 ModelConfig。这件事看起来平常，但我觉得很重要：dim、n_layers、n_heads、n_kv_heads、vocab_size、max_seq_len 这些参数如果散落在脚本里，后面做训练、推理、导出和复现实验都会乱。把它们收束到配置类里，模型才真正变成一个可调组件，而不是一坨写死的代码。

在结构上，这个小 LLM 仍然是典型的 Decoder-only 路线，但实现上有几个点很值得记住。RMSNorm 放在残差前面，稳定训练又不引入多余的均值计算；RoPE 把位置信息揉进注意力的几何关系里，不需要单独做位置 embedding；GQA 则通过较少的 KV 头去配更多的 Query 头，省显存，也更接近现在高效推理的思路。再加上 SwiGLU 风格的 MLP 和权重共享，整个模型会显得很“紧”，没有多余装饰。

我自己的理解是：这些设计不是为了让代码看起来更像 LLaMA，而是为了让小模型在有限资源下尽量保住表达能力。小模型最怕的不是少几层，而是每个模块都在悄悄浪费容量。

Tokenizer 不是附录

第二个关键点是 train_tokenizer.py。很多人会把 tokenizer 当成预处理脚本，但它其实决定了模型的语言边界。这里用的是 BPE，加上 ByteLevel 和 NFKC 规范化，目标很明确：既要兼容中文和英文混排，也要尽量减少未登录词带来的损失。

更关键的是特殊 token 的约定。<unk>、<s>、</s>、<|im_start|>、<|im_end|> 不是装饰品，它们直接影响模型怎么理解对话边界、起止边界和未知符号。chat_template 也很说明问题：如果训练时的模板和推理时的模板不一致，模型就会出现一种很烦人的现象，看起来“会说话”，但总是答不到正确位置上。

所以我现在会把 tokenizer 看成模型的一部分，而不是模型外面的附件。词表怎么切、特殊 token 怎么排、聊天模板怎么写，最后都会回到一个问题：模型到底学到了什么输入格式。

数据和 mask 决定训练到底在学什么

这一章后半部分最打动我的，是 PretrainDataset 和 SFTDataset 的对比。表面上它们都在把文本变成 X、Y，但真正的差异在 loss_mask。预训练阶段，模型学习的是通用的下一个 token 预测，除了 padding 之外，几乎整段文本都在参与学习；SFT 阶段则只让 assistant 的回答部分进入 loss，提示词和上下文更多是条件，而不是要背诵的答案。

这个设计把一件事说得很清楚：预训练负责“会续写”，SFT 负责“会按要求回答”。两者用的是同一个模型骨架，但目标完全不同。以前我总觉得对话能力是模型变大后自然长出来的，现在看更像是数据组织把能力方向拽过来了。

工程上这也很实用。模型答不对时，先别急着改优化器，先看 tokenizer 是否一致、样本格式是否统一、mask 有没有把 prompt 误算进 loss。很多训练问题，本质上不是模型太弱，而是监督信号传错了地方。

生成阶段才是检验点

generate 方法也值得单独看。它不是简单地把 logits 做个采样就完事，而是要处理最后一个位置的输出、温度、top-k、停止 token，还要兼容左侧 padding 和更长上下文的截断。也就是说，生成阶段不是训练后的附属功能，而是整个系统是否真的可用的最后一关。

这也解释了为什么很多模型训练 loss 看起来不差，实际一生成就开始跑偏。训练时看到的输入格式、mask 规则、停止边界，只要和推理阶段有一点偏差，最终体验都会被放大。

我的工程判断

这一章给我的结论很直接：做一个小 LLM，先盯三件事。第一，骨架是不是清楚，尤其是 norm、attention、MLP、weight tying 这些基础件；第二，Tokenizer 和 prompt 模板是不是和训练数据对齐；第三，loss_mask 和 generate 是否把“该学的”和“该停的”处理对了。

如果把这三件事做好，小模型至少能先跑通，并且能解释清楚自己为什么能跑通。对我来说，这比一开始就追求更大的参数量更重要。模型不是从某个神奇公式里长出来的，而是从一整条流水线里被一点点拼出来的。