先把丑话说在前面：这篇文章不是来宣布「我训出了一个吊打某某的强力小模型」的。实际上，我现在连能拿出来直接用的权重都还没训完。
发这篇出来，主要是想聊聊我把一个中文小模型从洗数据、搭架构一直到强化学习（RL）完整手搓一遍的过程。中间每一步是怎么设计的、为什么要这么选，以及目前在显卡上跑出来的一些阶段性数据。代码和洗好的数据我都开源了，家里有一张 16GB 显存的显卡就能直接跟着跑。
为什么要自己重造轮子？
之前看 MiniMind 和 nanochat，它们其实已经证明了「一个人从零手搓小模型」这条路完全走得通，代码写得也很清爽。我在做 tinybrain 的时候，就想在这个基础上再往前折腾几步：
- 数据自己从头洗，不直接拿别人现成洗好的数据集；
- 模型结构跟进 2025--2026 年的新架构，不再局限于标准的 GPT-2 或 LLaMA 结构；
- 后训练（Post-training）把全套流程走通：SFT、领域专家、在线蒸馏、强化学习全塞进去；
- 以中文为核心，而且我按个人喜好，专门塞了不少古文和诗词进去。
一、数据：26 个来源，分了四个档位
预训练数据是我从 26 个公开来源一点点扒拉过来的：包括中文 FineWeb-Edu、中英文维基、百度百科、维基文库、殆知阁古籍、诗词库、鲁迅全集，以及代码（UltraData-Code、Stack v3）、数学（Nemotron 教材、FineMath、NuminaMath 等）、论文（arXiv、peS2o），还有一部分英文网页和古腾堡公版书。
先说清洗：这部分全是最枯燥的脏活。修乱码、统一标点、繁转简；把太短的、话都说不利索的、语种不对的、纯广告和车轱辘话太多的文档全扬了。长篇的书和古籍按章节切开。
这里特别提个隐私合规的坑：各种手机号、身份证、邮箱我都用正则替换成了占位符。但之前本来想把法院裁判文书也加进去，结果抽查时发现里面带了大量当事人的真实姓名、出生日期甚至家庭住址，脱敏太难做干净，为了省事和避嫌，这整个数据源最后被我直接忍痛砍掉了。
跨来源去重：我的做法是取每篇文档最长的 3 句话当指纹。如果一篇文档里有 2 句话在优先级更高的文档里出现过，就判定为重复。为什么不设成 1 句？因为像常见的法律条文、名人名言，在很多完全不相干的文章里都会被引用，撞一句就删的话误伤率太高了。优先级方面，精校过的来源（法律、诗词、古籍、维基）排前面，普通网页抓取的排后面。这一套下来，大概砍掉了四分之一的数据。
分档：我定了个配比（中文文章 37%、百科 15.5%、代码 12%、数学 10%、论文 8%、古文和公版书 6%、英文网页 5%、英文书 4%......），算出每个来源在各个档位该取多少，然后按质量打分从高往下截。最精华的进 nano，往下依次是 mini、base 和 full。
| 档位 | Token 数 | 文档篇数 |
|---|---|---|
| nano | 0.99B | 54 万 |
| mini | 4.96B | 348 万 |
| base | 20.19B | 1319 万 |
| full | 60.26B | 4298 万 |
这么切的好处是，越小的档位单篇质量越高，适合快速验证；大档位的价值则在于管饱。另外，每条数据我都保留了 source 和 license 字段------毕竟是个大杂烩数据集，没法用单一许可证概括（比如维基是 CC BY-SA，还有两个上游来源压根没写许可证，我就先老老实实标成了「待确认」）。
二、分词器（Tokenizer）
用的是常规的 SentencePiece BPE，词表大小 32768，在按配比抽样的 3500 万 token 上训出来的。没做 Normalization，开了 byte fallback，数字全部逐位拆开，保证解码出来和原文一字不差。
测了一下压缩率，平均一个 token 包含的字符数：现代中文文章是 1.70 字，古文只有 1.11 字，英文网页是 4.16 个字符，综合下来是 2.38。古文压缩率这么低完全在意料之中------文言文基本是一字一义，BPE 确实很难把它们合并成词。
三、模型架构：塞了四个新玩具
既然是自己玩，架构上我就没因循守旧，主要试了四个新东西：
1. mHC 多流残差
普通的残差连接就一条主干道：x ← x + F(x)。mHC（DeepSeek 在字节 Hyper-Connections 基础上改的那个版本）相当于把马路拓宽成了 4 条车道。每个子层配三组系数：「读」负责把 4 条流加权合成本层的输入，「写」决定把输出按什么比例加回每条流，「混」则让不同流之间互相交换信息。
这里最核心的设计是把「混」约束成了双随机矩阵（每行每列加起来都等于 1，靠 Sinkhorn 迭代算出来）。这样信号总量在层与层之间传递时既不会爆炸也不会衰减------原版 Hyper-Connections 在深层训练时容易崩，病根就在这。在 tinybrain 里我用的是静态版，系数不依赖输入，所有子层的系数可以一次性算完，几乎不增加额外耗时。
2. Engram 查表记忆
语言里有很多死知识，属于「看到前几个字就条件反射想起后几个字」，拿宝贵的矩阵乘法去硬算其实挺浪费的。Engram 的思路很直接：按「前 n−1 个词 + 当前词」做哈希查表（n 取 2 和 3，各配 8 个哈希头）。查出来的向量，一部分当做额外的输入源，另一部分在第 1 层（medium 规格里还加了中间一层）之前，通过门控注入到残差流里。
因为查表每次只取十几行，完全不吃算力，只占显存，所以这张表可以铺得很大：在我的 small 和 medium 配置里，光这张记忆表就占了（主干 + 记忆表）总参数量的 45%。
3. 3:1 滑动窗口注意力
每 4 层里面，前 3 层只看最近 1024 个 token（局部窗口），到了第 4 层才看全文（全局注意力）。在 CUDA 上直接用 PyTorch 的 FlexAttention 编成块稀疏掩码，只算窗口内那条对角线带子。其他地方就是现在的标配了：GQA、QK-Norm、RoPE。
4. 混合专家 MoE（仅限 large 规格）
第 1 层保留普通的稠密 FFN，后面的层全部换成「2 个共享专家 + 48 选 8 的路由专家（宽度是正常的四分之一）」。
定这个比例前，我专门扒了 20 个公开 MoE 模型的配置做对比：2024 年以后的 MoE，激活比例基本都压在 12.5% 以下；和我们体量相近的小号 MoE，大多也是用 32 到 64 个专家。负载均衡方面，我抄了 DeepSeek-V3 的作业，主要靠一个不参与梯度的动态偏置（bias）来硬调，只辅以一个权重极小（1e-4）的传统均衡 loss。
最后定下来的三档模型尺寸：
- small：总参数 1.67 亿（主干 6400 万）
- medium：总参数 4.99 亿（主干 2.33 亿）
- large (MoE) ：总参数 11.8 亿（每次激活 2.66 亿）
四、预训练与踩坑
优化器配置：矩阵参数全部上 Muon（动量之后跑 5 次 Newton--Schulz 迭代做正交化），Embedding、LM Head 和 Engram 记忆表还是老老实实用 Adam。
分享一个很容易被忽略的暗坑：交叉熵是会「塌缩」的 。
刚好看到 Google DeepMind 今年发的一篇论文（arXiv 2610.31642）提到这个问题：当模型对某个预测非常有把握时，标准实现里的正确 token 概率在 FP32 精度下会被直接舍入成 1.0。这会导致正确项的梯度变成 0，偏偏其他非正确词的梯度还是正的，整份梯度加起来就不再等于 0 了。解决办法是换成数学上等价的 softplus(−m) 写法（m 是正确词 logit 领先其他所有词的最大间隔），这样就不会因为浮点数精度把梯度吞掉了。改完之后，我还特地在训练日志里加了一列，专门盯着到底有多少比例的 token 跨过了这个精度崩溃的临界值。
单张 5070 Ti (16GB) 到底能跑多快？
能上的加速手段我基本都拧满了：Liger-Kernel 的融合交叉熵和 RMSNorm、FlexAttention、可选的 cuDNN 滑动窗口注意力、逐层 torch.compile，再加上选择性激活重计算（Activation Checkpointing）。
实测下来，跑 small 规格（seq_len 2048 × batch 16），吞吐量大概在 7.1 万 token/秒 ，显存峰值 12.5 GiB；跑 medium （micro_batch 4，梯度累积 4），大概是 2.6 万 token/秒。
目前跑出来的一些中间数字（用的都是 20B 的 base 数据集，同一套验证集）：
- small 训到大概 25 亿 token 时，val loss 降到 2.88，目前我先暂停了；
- medium 在吃掉同样多 token 的情况下，loss 比 small 低了大概 0.17，而且随着训练推进，这个身位优势还在拉大（早期差距只有 0.13 左右）；
- 但现实问题是，medium 每一步要比 small 慢 2.6 倍。所以如果按**实际挂机时间（Wall-clock time）**来算，训到 4.4 小时左右，俩模型的 loss 才刚刚打平（3.035 vs 3.036）。不过按拟合出来的 scaling 曲线看，过了这个交叉点之后，medium 就会开始持续领跑了。
当然，这些都还是正在跑的中间状态，等彻底跑完，我会把最终的 loss 曲线和实际对话效果放出来。
五、后训练：麻雀虽小，五脏俱全
虽然模型不大，但后训练的整套活儿我也全给它安排上了：
- SFT：汇总了 18 个开源指令集（smoltalk-chinese、COIG-CQIA、Infinity-Instruct、OpenCoder 等），外加我自己攒的一批文学数据（命题写诗、文言白话互译、模仿鲁迅文风续写），只对回答部分算 loss。这里顺手洗掉了 400 多条大模型「自报家门」（比如自称我是某某公司的 AI）的脏数据。另外分了档：small 模型容量有限，直接去掉了思维链推理过程，长度截断到 2048；medium 和 large 则完整保留了 标签。
- 领域专家：拿通用 SFT 训出来的模型复制三份，分别在数学、代码、文学三个垂类数据上继续微调，训出三个偏科生。
- MOPD 多老师在线蒸馏：让学生模型自己生成回答，然后拉出对应领域的专家模型，在每个 token 位置上给出概率分布，去最小化反向 KL 散度。因为学生和老师用的是同一个分词器，所以这个 KL 散度是可以在整个词表上精确计算的。
- 强化学习（RL） ：在 Dr. GRPO 的基础上，我改成了每个 batch 只更新一次（完全 On-policy）。这种设定下，PPO 那个 clip 机制其实就失效了，损失函数直接退化成带组内基线（group baseline）的普通策略梯度。
为了给 RL 提供奖励信号，我写了 10 种题型的判分器（Verifier）：数学题直接对答案、代码题丢进沙箱跑单元测试、指令跟随用代码硬性检查约束、阅读理解算 F1、还有诗词填空和工具调用。最折腾的是诗词格律判分器：我写了个规则去按「一三五不论，二四六分明」硬抠平仄和押韵，然后再叠加上一个裁判模型（Qwen3.8-27B）给意境打分。
为了验证代码没有 bug，我先弄了个 2300 万参数的袖珍玩具模型，把「预训练 → 退火 → SFT → Agent RL → RLAIF → 专家微调 → RLVR → MOPD → 评测」这整条超长流水线在我的 Mac (M4 Max) 上从头到尾跑通了一遍，总共花了 198 分钟。当然，23M 这个体量的模型训出来基本是个复读机，评测分数跟瞎蒙差不多，跑这一遍纯粹是为了证明各个阶段的接口和梯度流向是正常的。
六、想自己动手试试？
如果你手里也有张 16G 的卡，直接敲下面几行命令就能跑起来：
Bash
bash
git clone https://github.com/cai114514114514/tinybrain && cd tinybrain
pip install -e ".[data]"
modelscope download --dataset caiwangzsnm/tinybrain-dataset pretrain_nano.parquet posttrain/sft_small.parquet --local_dir data
mv data/posttrain/* data/
tinybrain pretrain # 跑 nano 数据集，5070 Ti 上大概 4 个小时就能搞定
tinybrain sft
tinybrain chat接下来还要填的坑
目前这个项目还在继续肝，后面几周主要做这几件事：
- 把 small 和 medium 彻底训完，放出权重和正经的评测集跑分；
- 开训 large（MoE）版本；
- 补齐古文和诗词的专项评测；另外打算把诗词格律判分器里的字表换成真正的《平水韵》------现在图省事是用现代普通话拼音去近似的，遇到入声字经常判错，懂诗词的朋友看了可能会血压升高。
如果你最近也在自己折腾小模型，非常欢迎在评论区聊聊，尤其是关于数据配比怎么调 、小模型用什么新结构划算这块的实战经验。如果觉得这套代码和数据对你有点参考价值，去 GitHub 帮我点个 Star 就是最大的鼓励了！