ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

从零构建小型语言模型:手写Transformer核心模块实战

从零构建小型语言模型:手写Transformer核心模块实战 两年前我建了一个名为ai-engineering-from-scratch的项目目录打算做一件看起来有点“反效率”的事不调用任何封装好的大模型 API也不直接用现成的 Transformer 库而是从矩阵乘法开始一步步搭出一个小型语言模型。后来这个项目成了我理解深度学习最扎实的一段经历。先解释一下标题里的 Scratch它和少儿编程平台 Scratch 没有关系这里用的是英文短语 from scratch意思是白手起家、从零开始。如果你已经会写 Python但翻开大模型源码时经常一脸懵或者你跟着教程把模型跑通却说不清每个模块在做什么那这篇文章应该能给你一条完整的从零路径。我会讲清楚怎么选型、怎么搭数据管道、手写 Transformer 核心模块、处理训练中的各类翻车现场最后再把模型扩成工程化思路。1. 项目全貌为什么值得从零开始做AI工程1.1 是“从零构建”而不是“调包”现在大部分人做 AI 项目第一步就是pip install transformers然后调用一行代码跑出结果。确实方便但问题在于一旦模型效果不好很少人能定位到是数据问题、模型结构问题、训练策略问题还是部署环境问题。我做ai-engineering-from-scratch这个项目就是想打破这种黑盒状态。所谓 from scratch不是要求你从电子线路重新发明计算机而是不跳过关键环节词汇编码、Embedding、多头注意力、残差结构、反向传播、训练循环每一步都自己写出来。做完之后会有一种“每个细节都受我控制”的感觉这种掌控感在排查问题时非常值钱。用现实生活类比一下调包调参有点像开一辆自动挡汽车你会踩油门和刹车就能上路从零搭模型像学手动挡你必须理解离合器怎么配合发动机转速。虽然日常开车可能用不上手动挡但当你遇到复杂路况时那种对机械原理的理解能帮你准确判断故障。AI 工程也是一样模型 fail 的时候框架不会告诉你真正原因只有你自己知道内部发生了什么。我也见过一些极端做法要求从零写自动微分、从零写 GPU 算子甚至从零写内存分配器。我觉得没必要。我们的核心目标是理解 Transformer 和 AI 工程流程而不是重新发明框架。所以这个项目里的“从零”是有边界的模型算法、网络结构、训练流程自己实现但张量计算和自动微分基础可以借助成熟框架。和“学开车不需要先造发动机”一个道理但你得知道发动机有哪些部件、它们怎么配合。1.2 这个项目适合谁以及它的边界这个项目适合三类人一是刚学完深度学习基础理论、想动手验证原理的开发者二是已经在用大模型 API但想深入底层、提升调试能力的一线工程师三是准备求职 AI 工程岗位、需要拿出一个能讲清楚细节的作品的人。不适合只想快速上线聊天机器人的业务团队因为从零搭模型的成本远高于直接用现成方案而且初期效果大概率不如成熟开源模型。项目边界要提前立好不会涉及分布式训练、生产级部署、大规模数据清洗也不会从零实现 GPU 算子。目标是在一个普通笔记本的 CPU 环境下用几万条文本训练一个微型语言模型让它能生成语法通顺、主题基本连贯的短文本。这个边界非常重要否则项目会无限膨胀。每当你被“要不加个 MoE”“要不要做 RLHF”这类念头干扰时提醒自己先跑通最小闭环其他以后再说。1.3 澄清误区from scratch 不是 Scratch标题发出来之后经常被搜索到“scratch编程小游戏”“Scratch 植物大战僵尸”这类词。这里必须澄清一下Scratch 是麻省理工学院开发的图形化编程平台用户拖拽积木块就能控制角色运动非常适合少儿编程启蒙。而标题里的 from scratch 是英文固定搭配意思是“从零开始、白手起家”和那个猫形吉祥物没有关系。如果你是想找 Scratch 课程这篇文章帮不到你但如果你想做“从零开始的 AI 工程”这里是了。不过换个角度看Scratch 对编程思维启蒙确实有贡献它让人第一次意识到“程序 指令的有序组合”。AI 工程也是同理尤其在这个 transformer 权重可以无限堆叠的时代理解“指令如何组合”反而比记忆某个 API 更重要。从这个意义上无论用图形化编程还是 Python 手写张量核心都是在训练一种拆解问题、构建系统的能力这也是我把这个项目命名为 from scratch 的原因。2. 核心设计与技术选型2.1 硬件与软件栈的取舍最开始我在一台没有 NVIDIA GPU 的 MacBook 上做实验。为了能在 CPU 上快速迭代模型必须做得足够小参数量控制在几百万级别而不是大模型。硬件限制其实会逼你做出正确的工程决策因为机器慢你会本能地减少无用计算、重视数据批处理效率、优化内存访问。如果你有 GPU当然更好可以把模型稍微放大一些但不要一上来就追求“大”。先在一个能全天跑回归测试的规模上把代码逻辑验证清楚再考虑扩展。软件栈方面我的选择是 Python 3.9 NumPy PyTorch。PyTorch 只作为张量计算和自动微分引擎不直接调用torch.nn.Transformer、torch.nn.MultiheadAttention这类高层封装。可能有人会问为什么不用纯 NumPy 手写所有反向传播那样更有“原教旨主义”味道但会让整个项目陷入梯度推导的泥潭偏离理解 Transformer 结构这个核心目标。下面用一张表对比三条路线方案实现难度理解深度推荐程度纯 Python 双层循环很高最深不适合完整模型只适合学基础NumPy 手写反向传播极高很深适合 micrograd 级别的小玩具PyTorch 张量 自动微分中等高本项目的推荐路线我选了第三条因为这样既能亲手实现 Embedding、Attention、LayerNorm、FeedForward又不用为了softmax的梯度矩阵推导浪费一个周末。真实 AI 工程里也是同样的逻辑框架只是计算引擎模型架构、数据管道、训练策略都掌握在自己手里。2.2 数据准备与预处理细节从零开始做 AI 工程数据准备往往比模型代码更花时间。我选了一个公开的古诗词数据集大概几万首诗先抽取其中 5000 首做小样本验证。预处理分三步转成简体、过滤过短或过长文本、统一标点。这种清洗看起来简单但对训练稳定性的影响远超想象。例如繁体字会让字符表膨胀半角全角标点会让同一个字符被拆成两个词。接下来是分词方案。我采用字符级 tokenizer因为中文本身以单字为基本单位字符级处理直接而且不需要分词库。对于英文场景字符级会拉长序列长度更推荐研究 BPE但 BPE 的合并规则实现复杂度会高不少。字符映射需要建立stoi字符到 id和itosid 到字符并加入pad、bos、eos、unk四个特殊 token。字符集一旦固定训练数据里没出现过的字符只能映射成unk所以数据清洗时最好把所有特殊符号都提前处理掉。序列构造也是一门小学问。我设置上下文长度context_len128每次截取前 128 个字符作为输入第 129 个字符作为目标然后按步长 64 滑动生成下一条样本。这样相邻样本之间有重叠可以最大化利用数据但也意味着模型会多次看到近似的内容。如果步长太小样本重叠过多模型容易“背答案”而不是学规律步长太大样本数会变少。通常保留一半左右的重叠是比较好的折中。2.3 评估指标与边界确立训练语言模型最核心的指标是交叉熵损失它衡量模型预测分布与真实 token 分布的差距。比如某个位置真实字符是“春”模型给出的概率是 0.1交叉熵就会比较大如果模型给出的概率是 0.9损失就小。为了更直观通常还会计算困惑度perplexity缩写 PPL它像是“模型在下一个字符选择上的平均犹豫程度”。随机猜测模型的 PPL 大约等于词表大小所以看到 PPL 从几千降到几十就能感知到模型在逐渐学到语言规律。项目边界在这里同样要有明确定义这个模型不需要像 GPT-4 那样生成惊艳内容也不必支持对话、总结等复杂指令。目标只有一个——在验证集上看到持续下降的损失并且生成文本中能出现常见的词序和搭配。有读者可能觉得这个目标太低但 AI 工程讲究“先跑通再优化”尤其是从零开始的探索一次成功的完整闭环比十个半成品有价值得多。3. 实操过程从零搭一个小型语言模型3.1 搭建数据加载管道数据加载管道把所有文本变成可以喂给模型的 batch。我用torch.utils.data.Dataset和DataLoader实现核心代码很少但要注意形状对齐import torch from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, data, seq_len, step): self.samples [] for i in range(0, len(data) - seq_len - 1, step): x data[i:iseq_len] y data[i1:iseq_len1] self.samples.append((torch.tensor(x), torch.tensor(y))) def __len__(self): return len(self.samples) def __getitem__(self, idx): return self.samples[idx]这里的x是长度为 128 的输入序列y是同样长度的目标序列只是整体右移一位。也就是说给定前 128 个字符监督信息是第 2 到第 129 个字符。通过DataLoader的batch_size32每个 batch 就变成一个形状为[32, 128]的输入张量和一个同样形状的目标张量。还有一个容易忽略的点DataLoader会把每个 batch 自动堆叠但不同样本的文本长度在这里天然一致因为都是从固定长度切出来的。如果你的序列长度不固定就需要自己实现 collate 函数做 padding 和 mask。我建议从固定长度开始否则你还没见到 attention 的威力就先把 padding mask 的 bug 调到怀疑人生。3.2 手写Embedding与位置编码Embedding 层的作用是把 token id 映射成稠密向量本质上就是一张可查询的权重表。在 PyTorch 里一行代码就可以定义token_embedding nn.Embedding(vocab_size, embed_dim)。不过为了不彻底“调包”我们得知道它背后做的是查表输入 token id 是[B, T]输出是[B, T, embed_dim]。但光有 Embedding 还不够因为注意力机制本身对位置不敏感。如果两个句子词序打乱注意力计算出来的结果几乎一样这显然不符合语言规律。所以需要加入位置编码。我采用经典的正弦余弦函数sinusoidal公式如下PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))直观理解不同位置拥有不同的相位信号让模型能感知“这是第 5 个词”和“这是第 10 个词”。代码实现也不复杂import math import torch def get_position_encoding(max_len, d_model): pe torch.zeros(max_len, d_model) pos torch.arange(0, max_len).unsqueeze(1).float() div torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(pos * div) pe[:, 1::2] torch.cos(pos * div) return pe然后在模型 forward 里把 token embedding 和位置编码相加得到最终输入。要不要用可学习位置编码也可以不少现代模型采用可学习版本在固定长度内效果更好。但我选 sinusoid 有一个实际理由它能直接外推训练时用 128 长度生成时扩展到 256也不需要重新训练。3.3 实现多头注意力与残差结构这是整个项目的核心。我先实现单头缩放点乘注意力把公式写清楚def scaled_dot_product_attention(q, k, v, maskNone): d_k q.size(-1) scores q k.transpose(-2, -1) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) return attn v, attn为什么要除以sqrt(d_k)这要从 softmax 的行为说起。如果q和k是标准正态分布它们点积的方差会随着维度d_k线性增长。维度很大时点积数值可能到几十甚至上百softmax 被推到极端饱和区梯度会大幅度变小训练几乎停滞。除以sqrt(d_k)可以把方差拉回到 1 附近让 softmax 保持在有区分度的区域。这是 transformer 论文里的一个小细节但直接影响训练稳定性。然后是多头注意力。把模型维度embed_dim分成num_heads个头每个头独立做注意力之后拼接起来再过一层线性投影。这样模型能在不同子空间同时关注不同类型的关系比如一个头关注句法关系另一个头关注词义搭配。代码实现class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() assert embed_dim % num_heads 0 self.head_dim embed_dim // num_heads self.num_heads num_heads self.qkv nn.Linear(embed_dim, 3 * embed_dim) self.proj nn.Linear(embed_dim, embed_dim) def forward(self, x, maskNone): B, T, C x.shape qkv self.qkv(x).reshape(B, T, self.num_heads, 3 * self.head_dim) q, k, v qkv.chunk(3, dim-1) q q.transpose(1, 2) k k.transpose(1, 2) v v.transpose(1, 2) attn_out, _ scaled_dot_product_attention(q, k, v, mask) attn_out attn_out.transpose(1, 2).contiguous().view(B, T, C) return self.proj(attn_out)这里有一个关键点语言模型是自回归的预测第t1个 token 时只能看到前t个 token绝对不能看到未来信息。所以在训练时需要给注意力矩阵加一个因果 maskcausal mask。具体做法是把未来位置的分数设为-inf这样 softmax 之后这些位置的概率变成 0。如果你不放心可以生成一个上三角全 1 的 mask 矩阵用来屏蔽对角线之后的位置。TransformerBlock 除了多头注意力还有 FeedForward、残差连接和 LayerNorm。我采用 Pre-LN 结构先做 LayerNorm再做 Attention 或 FeedForward最后把输入加回来作为残差。和 Post-LN 相比Pre-LN 在大模型中训练更稳定不容易出现梯度爆炸。FeedForward 基本结构是两层线性加一个激活函数我选 GELUclass FeedForward(nn.Module): def __init__(self, embed_dim, ff_dim): super().__init__() self.l1 nn.Linear(embed_dim, ff_dim) self.l2 nn.Linear(ff_dim, embed_dim) self.act nn.GELU() def forward(self, x): return self.l2(self.act(self.l1(x)))残差连接的作用是让梯度能够跳过一两个子层直接回传缓解深层网络中的信息衰减问题。LayerNorm 则是把每个位置的向量归一化到均值为 0、方差为 1保证训练平稳。3.4 构建完整语言模型把上述模块拼成一个小型语言模型我给它取名TinyLM。类结构大概如下class TinyLM(nn.Module): def __init__(self, vocab_size, embed_dim, num_heads, num_layers): super().__init__() self.token_embedding nn.Embedding(vocab_size, embed_dim) self.blocks nn.ModuleList([ TransformerBlock(embed_dim, num_heads) for _ in range(num_layers) ]) self.ln_f nn.LayerNorm(embed_dim) self.lm_head nn.Linear(embed_dim, vocab_size, biasFalse) def forward(self, idx): B, T idx.shape x self.token_embedding(idx) x x get_position_encoding(T, self.token_embedding.embedding_dim) for block in self.blocks: x block(x, causal_mask) return self.lm_head(self.ln_f(x))这里的causal_mask在训练前根据context_len生成一次形状为[T, T]避免每个 batch 重复生成。lm_head把最终的隐藏向量映射到词表大小得到每个位置上下一个 token 的 logits。输入idx形状是[B, T]输出形状是[B, T, vocab_size]。有一个可以优化的细节lm_head的权重和token_embedding的权重形状相同都是[vocab_size, embed_dim]。可以让它们共享同一份权重这样既减少参数又相当于给输出层施加了一种正则化实践中常能加快收敛。但并不是所有任务都适合权重绑定如果模型需要更强的表达能力分开效果可能更好。我建议两个版本都跑一次根据验证集损失决定。3.5 训练循环与文本生成训练超参数我设置如下embed_dim128num_heads4num_layers2context_len128batch_size32learning_rate3e-4。优化器用 Adam损失函数用 PyTorch 的CrossEntropyLoss。因为模型输出是[B, T, vocab_size]而目标是[B, T]需要先把输出 reshape 成[B*T, vocab_size]目标 reshape 成[B*T]再计算损失。训练循环就五步前向算 logits算 loss反向传播更新参数清零梯度。看起来很短但要注意每个 batch 送入模型之前都要确保输入张量是 long 类型而不是 float。我有一个朋友初次跑模型把 token 索引存成了浮点数组Embedding 层直接报错排查了很久才发现是 dtype 问题。生成文本时采用自回归方式给定初始上下文每次预测最右端位置的下一个 token再把它追加到序列末尾循环往复。这里不能直接取argmax因为贪心生成很容易陷入重复。我加了一个 temperature 参数和 top-k 采样def generate(model, start_tokens, max_new_tokens, temperature0.8, top_k40): model.eval() idx start_tokens for _ in range(max_new_tokens): logits model(idx)[:, -1, :] / temperature probs torch.softmax(logits, dim-1) if top_k is not None: values, _ torch.topk(probs, top_k) probs[probs values[:, -1:]] 0 probs probs / probs.sum(dim-1, keepdimTrue) next_token torch.multinomial(probs, num_samples1) idx torch.cat([idx, next_token], dim-1) return idxtemperature大于 1 会让分布更平生成多样性提高但也更容易胡说八道小于 1 则更保守。top_k40的意思是只保留概率最高的 40 个候选再重新归一化避免采样到概率极低的满嘴跑火车 token。这种方式对小模型特别友好能在创意和稳定性之间取得平衡。3.6 训练结果与评估案例我在这套配置下训练了不到一万步验证集损失从初始 8 左右逐步降到 2 上下。困惑度也明显下降说明模型开始学到字符间的概率关系。输入“春眠不觉晓”做生成时模型给出的续写虽然不能和训练集里的原句完全对上但已经能出现“春”“晓”“风”这类高频词的搭配句子读起来开始有语感。如果你用更大的数据集、更深一点的网络效果会更明显。不要迷信这几个数字不同随机种子、不同数据集清洗方式都会影响结果。真正的收获不是最终 loss 有多低而是你能看着 loss 曲线下降并且能解释它为什么下降。这时候你才算把 AI 工程从“调参数”变成了“调系统”。4. 常见问题与排查技巧实录4.1 损失不下降学习率与初始化最常见的问题是损失曲线基本不走或者震荡。先检查数据是不是x和y错位错了是不是 mask 把该看到的 token 也屏蔽了如果数据没问题再检查学习率。学习率太高会让 loss 来回跳太低会让更新几乎无效。我习惯先用 3e-4 试跑 200 步观察 loss 下降速度不下降就降一个数量级到 3e-5震荡就升到 1e-3。新手最容易忽略的是初始化。PyTorch 默认初始化大多没问题但只要自己实现了线性层或 Embedding最好用nn.init.xavier_uniform_或标准差为1/sqrt(fan_in)的初始化。初始化太小会让输出方差过小初始化太大会让 softmax 饱和两层都会让反向传播梯度消失或爆炸。4.2 注意力矩阵溢出与NaN训练中途 loss 突然变成 NaN是新手最崩溃的场景。绝大多数时候是数值稳定性问题。比如点积没有除以sqrt(d_k)大维度下 scores 变得很大softmax 指数溢出产生inf随后 loss 变成 NaN。另外一个被忽视的点是因果 mask如果某一行 mask 全部是 0softmax 会收到一列-inf结果就是 NaN。我在实现时习惯打印 mask 矩阵确认每一行至少有一个位置没有被 mask 掉。如果你用 float16 混合精度训练溢出概率会更高小模型建议先跑 float32稳定之后再切混合精度。4.3 过拟合模型容量与数据规模小模型更容易出现过拟合。我在 5000 首文本上训练到第二个 epoch 时训练 loss 还在降验证 loss 却开始回升。这种迹象说明模型已经记住了训练集中的片段。解决过拟合的第一选择是增加数据把 5000 首扩展到几万甚至几十万如果暂时没有更多数据就在 Attention 和 FeedForward 后面加 dropout通常设dropout0.1。也可以用早停每轮保存验证 loss 最低的模型不要直接用最后一个 epoch 的 checkpoint。训练时还有一个简单但实用的技巧做一个“最频繁字符”基线模型预测每个位置最常出现的字符看它的 loss 是多少。如果你的模型 loss 连这个基线都远不如说明模型根本没开始学有效特征问题大概率出现在数据管道而不是模型复杂度。4.4 生成结果重复解码策略模型训练完了生成出来的文本却不断重复这种问题几乎人人都遇到。比如输入“大漠孤烟直”之后模型一直输出“大漠孤烟直大漠孤烟直”这其实不代表模型没学好而是采样策略太单一。用argmax贪心解码会固定走概率最高的路径很容易掉进重复循环。解决办法是引入随机性设置temperature0.8从 softmax 分布里采样再叠加top_k40只保留概率最高的 40 个候选重新归一化。如果还不行可以加惩罚在每一步计算时对已经生成过的 token 的 logits 减去一个小常数让模型更不愿意重复自己。这个 n-gram 惩罚在文本生成库里很常见手写起来也不难。4.5 工程化问题代码组织与可复现性当项目从一个脚本膨胀到多个文件时代码组织不好会拖慢所有实验。我的目录结构很简单data.py放数据加载和预处理model.py放模型结构train.py放训练循环generate.py放推理脚本config.py集中放超参数。所有随机种子都要固定random.seed(42)、np.random.seed(42)、torch.manual_seed(42)。模型 checkpoint 里不光要存state_dict还要把超参数和 tokenizer 映射一起存成 JSON。否则过两周你再打开项目很可能想不起当时的词表长什么样。我每次实验都往logs/exp1.csv写一行损失、学习率、时间和 git commit 号后面画图、对比实验都靠它。这些习惯很朴素但能救你于水火。5. 从手写模型到AI工程化扩展与反思5.1 从原型到产品的距离我常提醒自己这个微型语言模型只是个学习原型离生产级 AI 工程还有很长的路。生产环境里你需要处理分布式训练、数据版本管理、模型灰度上线、推理延迟优化、反馈闭环等一系列问题。以推理为例我这个生成函数每次都要重新跑一遍全序列的 forward没有做 KV Cache速度很慢真正服务线上请求时需要用 KV Cache 缓存历史 token 的 Key 和 Value才能实现增量解码。还有 batch 服务多个请求拼成一个 batch 同时推理能显著提升吞吐。这些工程细节只有你理解了 attention 内部的数据流动之后才能真的明白为什么要这样做。5.2 下一步的进阶路径如果还想继续深入我建议按这个顺序扩展。第一步把字符级 tokenizer 换成 BPE词表从几千扩展到几万编码效率和模型语义上限都会提升。第二步加入指令微调用一批“问题-答案”对继续训练模型让它学会跟随指令。第三步尝试更高效的 attention 变体比如 Flash Attention它通过分块计算和在线 softmax 减少显存占用是当前大模型训练的重要基础。最近也看到不少人在探索“build a reasoning model from scratch”路线想让模型具备推理能力但底层步骤还是离不开预训练、微调、对齐这条主线。我的建议是不要一口吃成胖子每个阶段都跑通一个端到端闭环再迈进下一步。5.3 几点实战心得最后分享几条个人感悟。第一动手前先写清楚“不做什么”否则项目会无限膨胀。第二每实现一个模块立刻用很小的测试用例验证形状和梯度比如构造一个[2, 4]的随机输入跑一遍 forward再检查输出维度对不对。第三保存模型和配置时务必同时保存当前 git 提交号否则实验结果无法复现。第四遇到问题不要急着改代码先记录当前现象和所有输入条件否则很难对比改前改后。第五从最小的可运行版本开始先把流程跑通再慢慢优化。这个过程最让我受益的不是最终生成效果多好而是建立了一种“我明白这行代码到底在做什么”的底气。
返回列表