技术札记

从随机权重到本地推理:我如何从零训练一个 Transformer

在一台 16GB MacBook Air 上,用 MLX 从头实现 Decoder-only Transformer,并经历 Byte Tokenizer、BPE、Unicode、安全推理和独立评测的三代实验。

训练一个 Transformer,最容易做的事情是调用现成框架,最容易跳过的事情则是理解一条文本如何变成 token、梯度如何改变随机权重、Checkpoint 如何重新加载,以及为什么一个验证 Loss 很低的模型仍然不会回答真实问题。

我建立了 Local Transformer Training Lab,目标不是做一个能与大模型竞争的助手,而是在个人电脑上亲手跑通:

数据生成
→ Tokenizer
→ Decoder-only Transformer
→ 预训练
→ 指令微调
→ 独立评估
→ Checkpoint 与导出
→ 本地推理

项目在一台 16GB 内存的 MacBook Air 上完成,使用 Apple MLX 和 Metal GPU。所有模型都从随机权重开始训练,不是对 Qwen、Llama 等现成模型做微调。

三代本地 Transformer 实验的演进

为什么要从零训练

调用模型 API 能让我快速构建产品,却很难回答一些更底层的问题:

从零训练的价值,不在于参数规模,而在于任何异常都不能推给“黑盒”。数据、模型、训练和推理之间的因果关系会变得具体。

第一代:先让完整链路跑起来

第一版 programmer-world-tiny 使用本地模板生成约 10 MiB“程序员世界”语料,覆盖 Go、Linux、Redis、MySQL、Docker、代码补全和招聘面试:

项目数值
总记录23,377
训练 / 验证22,908 / 469
文本重复率0
外部 API未使用

模型保持刻意的小:4 层 Decoder Block、4 个注意力头、128 隐藏维度和 384 FFN。128-token 版本有 742,272 个参数;上下文扩到 512 后有 791,424 个参数。

核心模块全部在 src/model.py 中实现:

Token Embedding + Position Embedding
→ RMSNorm
→ Causal Multi-Head Self-Attention
→ Residual
→ RMSNorm
→ GELU Feed Forward
→ Residual
→ LM Head

因果 Mask 保证位置 t 只能看到自己和之前的 token。训练目标没有神秘之处:输入序列去掉最后一个 token,标签序列去掉第一个 token,模型学习预测下一个 token。

logits = model(inputs)
loss = cross_entropy(logits, targets)

训练使用 AdamW、梯度裁剪、定期验证,并同时保存 best、阶段 Checkpoint 和 final。最终模型包只依赖三类文件:

config.json
model.safetensors
tokenizer.json

这三者缺一不可。权重没有配置无法恢复结构,模型没有同版本 Tokenizer 无法解释 token id。

Byte Tokenizer:简单、可靠,也很昂贵

第一版使用 UTF-8 字节级 Tokenizer,词表只有 263 个 token:256 个字节值加特殊标记。它不需要训练、可以无损覆盖中文英文和代码,非常适合理解流程。

代价也很明显:一个中文字符通常需要三个 token。上下文窗口标称 512,能装下的中文远少于 512 个字。模型还可能逐字节生成一个不完整的 UTF-8 序列,终端最终显示

第一代训练 2,500 步,512-token 版本最佳验证 Loss 为 0.135433,最终困惑度 1.1514,耗时约 110.63 秒。数字非常漂亮,但真实问答并不可靠。

这是项目最重要的早期提醒:训练集和验证集来自同一套模板分布,低 Loss 首先说明模型记住了模板规律,不等于获得了通用编程知识。

第二代:BPE、两阶段训练与回答区间 Loss

v1 把模型扩大到 8,534,272 参数,使用 8 层、8 头、256 隐藏维度,并训练 4096 词表的 ByteLevel BPE。

Tokenizer 的效率变化很直观:

“Redis 为什么这么快?”
Byte Tokenizer:28 token
BPE Tokenizer:7 token

更重要的变化是把训练拆成两个阶段。

领域预训练

预训练数据不再保留问答外壳,而被改写为技术笔记和完整代码。目标是学习领域语言、代码结构和 next-token 规律。

指令微调

SFT 数据清除模板化的“补充约束”“模拟业务背景”等表达,做语义去重和类别均衡。训练时只对 <|assistant|> 之后的回答区间计算 Loss,避免让模型把大量容量用在复现用户问题和控制标记上。

v1 完成 800 步预训练和 1,500 步 SFT。最低 SFT 验证 Loss 达到 0.235517,但独立改写提示暴露了真实问题:Redis、MySQL、Docker 之间会串台,身份回答会混入 HR 模板,换一种问法后泛化明显下降,少数提示直接生成结束标记。

这次实验比“Loss 降了多少”更有价值的结论是:

参数规模、BPE 和训练算法可以改善表示与拟合;语料的自然多样性、监督质量和独立评测决定模型是不是真的会回答。

第三代:先保证输出是合法的

v1 的 ByteLevel BPE 仍允许模型生成无法组成合法 UTF-8 的字节片段。v2 没有继续盲目扩大参数,而是换成 Unicode 字符 Tokenizer。

训练语料包含 750 个不同字符,加 9 个特殊标记,最终词表为 759。每个输出 token 都是合法 Unicode 码点;未知字符明确变成 <|unk|>,不再产生

v2 有 6,825,728 个参数,仍然是 8 层、8 头、512 上下文,完成 800 步预训练和 1,500 步 SFT,总训练耗时 115.45 秒。

字符级 Tokenizer 并不一定比 BPE 更强。它的序列更长,当前又没有 KV Cache,生成速度更慢。但它验证了一个工程原则:模型能力、Tokenizer 效率和输出合法性是三种不同目标,不能只用一个 Loss 概括。

推理不是训练脚本的附属品

一个能加载权重的脚本,不等于一个可用的推理入口。v2 在生成层增加了:

其中最反直觉的是默认关闭历史。拥有 512-token 上下文不代表模型学会了多轮对话。如果训练数据没有可靠的多轮结构,把上一次错误回答放进下一轮上下文,只会让错误复利。能力没有准备好时,少做比假装支持更可靠。

如何评估一个教学型小模型

训练和验证按同一模板随机切分,很容易产生过度乐观的指标。项目最终采用三层检查:

  1. 实现正确性:Tokenizer 往返、Mask、形状、保存加载和最小训练测试;
  2. 同分布指标:训练 Loss、验证 Loss、困惑度和梯度范数;
  3. 分布外行为:人工改写问题、领域串台、空回答、重复、乱码和多轮污染。

真正值得保留的不是某个最好看的数字,而是失败案例。它们指出下一步应该改数据、Tokenizer、Loss、模型还是推理策略。

我真正学到的五件事

1. 数据结构比数据体积更重要

10 MiB 高度模板化语料可以得到很低的 Loss,也可以得到很差的真实泛化。去重不仅要比较字符串,还要避免语义模板跨越训练和验证集合。

2. Tokenizer 是模型架构的一部分

它决定词表参数、序列长度、上下文利用率、未知字符行为和输出是否合法,不是一个可以随意替换的预处理工具。

3. 预训练与 SFT 解决不同问题

预训练学习语言和领域分布,SFT 学习如何按照任务格式回答。把两种数据混在一次 next-token 训练里,模型很难知道哪些内容是知识,哪些内容是交互协议。

4. Loss 只回答它被设计来回答的问题

同模板验证 Loss 回答的是“模型能否预测同分布 token”,而不是“它能否理解用户的自然表达”。指标必须与产品声称的能力对应。

5. 推理治理不是掩盖模型缺陷

重复停止、回退和关闭历史不会让模型更聪明,但会让能力边界诚实、稳定、可观察。可靠系统不只追求更多正确输出,也要限制错误输出如何扩散。

下一步

这个项目已经完成从随机权重到本地推理的闭环,但仍然是教学模型。下一阶段最有价值的工作不是继续重复合成模板或盲目增加训练步数,而是:

从零训练没有让我更轻视大模型,恰恰相反。它让我看到,一个看起来简单的回答背后,数据、Tokenizer、架构、优化、评测和推理策略必须同时成立。

仓库与完整实验记录:hh696-wq/local-transformer-training-lab

返回技术札记