从零手写Transformer文本分类全链路工程实践

写这篇东西的起因，是我最近把公司一个内部文本分类项目彻底翻新了一遍，核心目标只有一个：把整条链路——从原始文本到训练样本、从模型前向到反向传播、从训练循环到推理接口——全部用自己代码实现，不直接套现成的 Trainer、不无脑加载预训练全家桶，也就是所谓的ai-engineering from scratch。折腾完这轮，我对“工程”二字的理解比过去三年调参加起来都深。这篇文章把整个过程中的设计思路、关键实现、踩坑记录都整理出来，给想从“会用”走向“能造”的读者，尤其适合每天跟模型打交道但总感觉隔了一层黑盒的同学。

先说清楚一点：from scratch 不是让你所有东西都从零发明。底层张量运算、基础数学函数、文件 IO 这类基础设施，用 NumPy 和 PyTorch 没有任何问题。真正的“从零”指的是不要依赖别人封装好的高层抽象——模型结构自己搭、数据管道自己写、训练循环自己控、服务接口自己出。这个尺度拿捏好，你的学习收益和工程自由度是最大的。

1. 项目整体设计与思路拆解

1.1 为什么要把全链路改成 from scratch

以前我的工作流特别“标准”：拿 HuggingFace 的 Trainer 一拉，数据集丢进去，训练一跑，指标一打印，完事。确实快，但有个致命问题——模型出问题时你完全不知道去哪里排查。比如 loss 不收敛，到底是数据没对齐、学习率不合理，还是模型实现有 bug？Trainer 把 80% 的细节都封装掉了，留给你的只有超参数面板，你除了调 learning rate 和 batch size 之外几乎无事可做。

这次动手之前，我把过去半年在生产环境里遇到的诡异问题列了个清单，发现绝大多数都指向同一个根源：我对底层机制的理解是碎片化的。数据 padding 的 mask 到底是怎么进 attention 的？梯度累积的缩放到底该放在什么位置？warmup 和 weight decay 为什么组合起来会有奇怪的效果？这些问题我都能说个大概，但从来没亲手实现过一遍。

所以这个项目在设计阶段的目标就三条：

• 用最小但完整的实现，把数据流动、梯度流动、参数更新这三大循环在自己手里跑通。
• 暴露所有中间张量的形状和数值，让每一个可调参数都有迹可循。
• 保持代码体积可控，单文件能读懂，方便后续往里加自定义逻辑。

第一版工程我控制在三个文件加一个配置文件，一共不到两千行。这本身就是 from scratch 的一个重要收益：当你的代码量小到能全量读一遍，你对系统的掌控感会完全不同。

1.2 技术选型：哪些必须手写，哪些可以依赖底层库

很多人一想到 from scratch 就陷入“什么都要自己造轮子”的极端，比如连矩阵乘法都要手写循环，这是没有必要的。我这次的划分方式比较实用，列成一张表供你参考：

这个选型的核心思路是：框架负责“数学运算”，你负责“工程逻辑”。模型定义、数据流动、训练策略、状态管理全部自己写，底层只当计算器用。这样出了问题，你可以一行一行看自己的代码，而不需要钻进别人的封装里去猜。

1.3 整体架构与模块边界

整个项目的模块边界我画得很干脆，四条流水线互不干扰：

• 数据流水线：raw text → clean → tokenize → build vocab → encode → bucket padding → batch tensor。
• 模型流水线：embedding → transformer block × N → pool → classifier → logits。
• 训练流水线：forward → loss → backward → clip grad → optimizer step → lr schedule → checkpoint。
• 推理流水线：加载 config 和 checkpoint → 走同一个模型 forward → softmax → 输出 top-k。

模块边界清楚之后，最大的好处是可以独立测试。数据流水线产出什么形状的张量，模型 forward 消费什么形状的张量，这两边只要接口对齐，单独调任何一条都不影响另外的。我在写代码的过程中也一直保持这个习惯：每个模块先写一个最小的烟囱测试，通了再拼起来。

2. 核心细节解析与实操要点

2.1 数据层：从原始文件到训练样本的完整链路

数据层的核心不只是“读文件”，而是把原始文本变成模型能吃的整数张量的全过程。我这次用一个简化但完整的字符级 tokenizer 来承载整个构建过程，没有直接用 BERT 的分词器，原因很简单：字符级 tokenizer 的词表构建逻辑最透明，特别适合展示 from scratch 的思路。

第一步是清洗。原始文本里有大量 HTML 标签、URL、重复标点。我的清洗规则是：小写化、去 URL、去 HTML 实体、把连续空格压缩成一个。这些规则看着简单，但在文本分类场景里对指标影响极大。实测同一个模型，清洗前后准确率能差 5 个百分点以上，因为很多噪声字符直接干扰了词表分布。

第二步是构建词表。我对每个字符计数，按出现频率排序，保留前 5000 个字符，再加四个特殊 token：<pad>（填充）、<unk>（未知）、<cls>（分类）、<eos>（结束）。词表构建完，我把频率分布画了个图，发现长尾非常严重，前 200 个字符就占了 80% 的 token 数量——这提醒我在后面的 embedding 层要注意初始化范围，否则低频字符的 embedding 很难被训练到。

第三步是编码与 padding。编码就是 token 到 id 的映射，padding 是让一个 batch 内所有样本等长。这里有个经常被忽略的点：padding 方向。我特意让<pad>放在序列末尾，同时维护一个 attention mask 张量，标记哪些位置是真实 token、哪些位置是 pad。这个 mask 后面在 attention 计算中要用来屏蔽 pad 位置的信息传播，如果你直接忽略它，模型会学到把 padding 当成有效信息，性能直接打折。

第四步是构造 batch。我实现了长度相似的样本凑在一起的采样策略，而不是全随机采样。原理很简单：一个 batch 的 padding 量由最长样本决定。假设你有 100 条长度 50 的样本和 1 条长度 500 的样本，随机放一个 batch，整个 batch 都要 pad 到 500，GPU 显存大部分浪费了。长度分桶之后，短样本放在一起算，速度能提升两三倍。

2.2 模型层：手写一个微型 Transformer 的核心组件

模型层是这个项目的核心学习区，我带着读者把每个组件过一遍，从输入到输出全部自己搭。

首先是Embedding 层。我建了一个词表大小 $V=5002$、维度 $d=128$ 的查找表，初始化用均值为 0、标准差为 $0.02$ 的正态分布。这个初始化值很关键，太大了会导致后续 attention 的 softmax 结果过于尖锐，太小会梯度消失。我踩过把标准差设成 0.1 的坑，结果训练一开始 loss 就震荡得很厉害。

然后是多头注意力。这一步是最容易写错的地方。注意力的公式是：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + \text{mask}\right) V $$

我设了 4 个头，每个头维度是 32，四个头拼起来又回到 128 维。实操中我提醒自己注意三个细节：

• 缩放因子 $\sqrt{d_k}$ 不能省。如果不缩放，当 $d_k$ 是 32 时，点积结果的标准差大约在 5.6 左右，softmax 之后几乎变成了 one-hot，梯度很小，模型学得很慢。缩放之后，结果分布更平滑，梯度能正常流动。
• mask 要用大负数而不是 0。padding 位置在加 mask 时要加-1e9，这样经过 softmax 之后对应位置的权重接近 0。如果你用 0 去加，softmax 里 $e^0=1$，padding 位置的注意力权重依然非零，信息就泄漏了。我用这个小细节至少排查了一个下午。
• mask 的形状是 [batch, 1, seq_len]还是 [batch, 1, 1, seq_len]*，取决于有没有多头维度。我用 PyTorch 自动广播，写成[batch, 1, seq_len]，在多头维度上自动广播，省显存也省代码。

接下来是残差连接和 LayerNorm。Transformer Block 的标准结构是“注意力子层 + 残差 + LayerNorm”和“前馈子层 + 残差 + LayerNorm”。LayerNorm 做的是对最后一维做归一化，每个 token 单独归一，而不是对整个序列，这个维度区分很重要。我推荐把 LayerNorm 放在残差相加之后，即“Post-Norm”，这种结构在小模型从零训练时稳定性更好。

最后是前馈网络。就是一个两层的 MLP：$128 \to 512 \to 128$，中间用 GELU 激活。这里加了一层 Dropout，概率 0.1。Dropout 的位置踩过坑：放在激活函数之后而不是之前。放在之前会把激活函数的负值区域给丢掉，实际效果差很多。

模型最后接一个分类头：把序列第一个 token（即<cls>位置）的输出取出来，经过一个线性层映射到类别数。这个操作叫“CLS pooling”，原理是让模型在训练中学会把全局信息汇聚到第一个 token 上，用它的向量代表整个序列。

2.3 训练层：优化器、学习率调度与梯度裁剪的工程细节

训练层作为 from scratch 项目里最容易“跑通但跑不好”的部分，很多细节都藏在常见封装之下。这次我把每一步都显式写出来，终于理解了之前很多调参经验背后的原理。

优化器我用的是 AdamW，未直接用库里的torch.optim.AdamW一行调完，而把参数分组方式手动配置：所有 embedding 层的参数和注意力层的参数用不同 weight decay。具体来说，偏置项和 LayerNorm 的 gamma/beta 不做 weight decay，其余权重做 0.01 的 weight decay。这个处理原本是预训练大模型时的经验，但实测放到小模型上也有效，能略微提升泛化，并且不会把 LayerNorm 的缩放能力给限制住。

学习率调度我加了 warmup。具体做法是：训练总步数 10000 步，前 500 步是 warmup，学习率从 0 线性升到峰值 1e-3；后面 9500 步用余弦退火，从峰值逐渐降到峰值 10%。注意这里的梯度问题——warmup 不是“可选项”而是“保护选项”。如果不做 warmup，前几个 batch 的梯度方向可能和整体数据集方向差异很大，学习率过大会一步把 embedding 和 attention 的参数推到一个很差的位置，后面很难回来。我把初始学习率改成 3e-4、不做 warmup 对比过，loss 曲线前 200 步就开始震荡，之后无论怎么调都回不到正常水平。

梯度裁剪也是必需品。我设置的 max_grad_norm 是 1.0。梯度裁剪的原理不复杂：把整个参数梯度的全局范数限制在 1.0 以内，如果超过就等比例缩放。它不能替代正确的学习率，但能避免偶发样本造成的梯度尖峰把参数直接推出有效区域。这个小样本测试里，不加梯度裁剪跑了 2000 步出现过 3 次 loss 突然跳到正常值 10 倍以上的情况，加上裁剪后一次都没出现。

训练循环是另一个容易“看着会、写不对”的地方。梯度累积的实现方式我踩过一个坑：直接把 4 个 batch 的 loss 相加除以 4 当作平均 loss，其实应该在每个 micro batch 的 loss 上做除法，或者每次累加时除以累积步数。否则 loss 规模比真实值大好几倍，会连带影响梯度裁剪阈值判断的准确性。

训练过程中的周期性评估我也手动实现了一遍：每 500 步在验证集上跑一遍前向，记录准确率和 loss，存出最佳模型。这个逻辑本身简单，但有一个细节值得写进你的代码里——评估时一定要切到model.eval()模式并关闭梯度计算，否则 dropout 和 batch norm 的随机性会污染评估结果，同时显存也会被不必要的计算图占住。

3. 实操过程与核心环节实现

3.1 搭建可复现的数据加载模块

数据加载模块我一共写了三个类：TextDataset、LengthBucketSampler和Collator。核心目标是高度可复现。所谓可复现，就是同一条文本在任何时刻、任何进程里被处理出来，得到的 id 序列必须一模一样。这里必须手动指定随机种子，并且在采样器里也要固定，否则跨机器跑结果就对不上。

TextDataset做的事情是读入原始文件，应用清洗规则，再做 tokenizer 编码。我的代码如下：

import re from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, file_path, vocab, max_len=256): self.samples = [] self.max_len = max_len self.vocab = vocab self.pad_id = vocab["<pad>"] with open(file_path, encoding="utf-8") as f: for line in f: text, label = line.rstrip("\n").split("\t") ids = self._encode(self._clean(text))[:max_len] self.samples.append((ids, int(label))) def _clean(self, text): text = re.sub(r"http\S+", "", text) # 去 URL text = re.sub(r"<[^>]+>", "", text) # 去 HTML 标签 text = re.sub(r"\s+", " ", text).strip() # 压空格 return text.lower() def _encode(self, text): return [self.vocab.get(ch, self.vocab["<unk>"]) for ch in text] def __len__(self): return len(self.samples) def __getitem__(self, idx): ids, label = self.samples[idx] return ids, label

LengthBucketSampler的思路是按照长度把所有样本索引分组，每个 batch 内的样本长度尽量接近。这里有个小技巧：不需要精确排序，只需要按长度粗略分成几个桶，比如 16 个桶，再在桶内随机取样组成 batch。精度不重要，但效率提升很显著。

Collator负责把一个 batch 的样本列表变成两个张量：input_ids和attention_mask。代码重点是 padding 的逻辑：

def collate_fn(batch, pad_id=0): max_len = max(len(ids) for ids, _ in batch) input_ids, attention_mask, labels = [], [], [] for ids, label in batch: padded = ids + [pad_id] * (max_len - len(ids)) mask = [1] * len(ids) + [0] * (max_len - len(ids)) input_ids.append(padded) attention_mask.append(mask) labels.append(label) return torch.tensor(input_ids), torch.tensor(attention_mask), torch.tensor(labels)

注意我 padding 放的是0，也就是词表里<pad>对应的 id。attention mask 用 1 表示有效 token、0 表示 pad，后面在注意力计算里会用到。

3.2 手写前向传播与反向传播

模型部分我实现了一个两层的微型 Transformer，核心的TransformerBlock长这样：

import torch import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, d_model=128, n_heads=4, dropout=0.1): super().__init__() assert d_model % n_heads == 0 self.d_k = d_model // n_heads self.n_heads = n_heads self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_o = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): b, n, _ = x.shape q = self.w_q(x).view(b, n, self.n_heads, self.d_k).transpose(1, 2) k = self.w_k(x).view(b, n, self.n_heads, self.d_k).transpose(1, 2) v = self.w_v(x).view(b, n, self.n_heads, self.d_k).transpose(1, 2) scores = q @ k.transpose(-2, -1) / (self.d_k ** 0.5) if mask is not None: mask = mask.unsqueeze(1).unsqueeze(2) scores = scores.masked_fill(mask == 0, -1e9) attn = self.dropout(torch.softmax(scores, dim=-1)) out = (attn @ v).transpose(1, 2).contiguous().view(b, n, -1) return self.w_o(out) class TransformerBlock(nn.Module): def __init__(self, d_model=128, n_heads=4, d_ff=512, dropout=0.1): super().__init__() self.attn = MultiHeadAttention(d_model, n_heads, dropout) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) def forward(self, x, mask=None): x = x + self.attn(self.norm1(x), mask) x = x + self.ffn(self.norm2(x)) return x

关于反向传播，直接使用 PyTorch 的 autograd。关于“前向都手写了，反向也要手写吗”这个问题，我的回答是：不需要。你手写复杂模型的反向只是为了证明你懂链式法则，对你的工程能力提升有限，还容易引入梯度错误。真正的工程价值在于你理解反向传播是由前向的计算图自动构建的，所以你在前向里的每一个操作都能影响梯度流向。只要保证前向是完全可微的，反向就交给 autograd。

但也有个地方值得手动看一下梯度：用loss.backward()之后，打印几层参数的梯度范数，确认梯度没有消失也没有爆炸。这是我在训练小模型时最常用的诊断手段。

3.3 训练循环与 checkpoint 管理

训练循环是整条链路的节奏控制中心。我直接贴一个模块化训练循环的核心代码：

for step, batch in enumerate(train_loader): input_ids, attention_mask, labels = [x.to(device) for x in batch] logits = model(input_ids, attention_mask) loss = criterion(logits, labels) / accum_steps loss.backward() if (step + 1) % accum_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() lr_scheduler.step() optimizer.zero_grad()

这里有个细节：criterion(logits, labels) / accum_steps放在 forward 之后而不是 backward 之前，是为了让梯度幅度和单 batch 训练保持一致。如果你直接用累加的 loss 去 backward，梯度范数会变成原来的accum_steps倍，那么后面梯度裁剪的阈值判断全部失效。

Checkpoint 管理我做了两件事：第一，每 500 步评估一次验证集；第二，只有当验证准确率创新高时才保存模型。保存的内容包括模型参数state_dict、优化器状态、当前学习率、step 数，方便中断之后从 checkpoint 恢复继续训练。很多人只保存模型权重，中断之后要从头再来，白白浪费算力。优化器状态不保存的话，即使模型参数恢复到最佳状态，学习率调度和动量信息也是断的。

4. 常见问题与排查技巧实录

4.1 loss 不降 / 反复横跳的排查思路

loss 不降是 from scratch 训练里最常见的噩梦，而且 90% 的情况下不是模型结构的问题，而是数据或配置的问题。我的排查顺序固定如下：

• 先确认数据对齐。取一个 batch 出来，手算一下真实 label 分布，再打印一下 logits 的数值范围。如果 logits 全是 0.001 这种数值，初始化可能太小；如果全是几百上千，初始化太大。
• 再确认 tokenizer。随机抽几条样本，把 id 还原成字符，看看是不是清洗规则把有效信息都删了。我遇到过因为标点清理太猛，把句子里的中文逗号全删了，导致语义翻转的诡异情况。
• 接着看学习率。如果你设的 peak 学习率超过 1e-3，先降到 1e-4 试跑 200 步，看 loss 有没有变化。这不是在给你推荐保守配置，而是为了判断是不是学习率过大导致震荡。我调试时通常先跑小学习率确认训练方向没问题，再逐步放大。
• 最后才查模型内部。打印每个 transformer block 的输出标准差。如果某层输出全部变成同一个值，说明该层可能退化了，常见原因是 LayerNorm 作用和维度不对。

4.2 梯度爆炸与 NaN 的定位

梯度爆炸最典型的特征是 loss 突然出现 inf 或 nan，好嘛一下子回到解放前。我的排查方式是用 hook 在参数上挂一个梯度范数监控器，打印每个参数组的梯度范数直方图。我知道问题发生时是哪一层先爆的，这比盲目调整学习率高效得多。有一次训练 nan 的罪魁祸首竟然是 attention 分数里出现了一个特别大的正值，softmax 之后对应的行梯度变成 nan。我用torch.isnan(scores).any()在内部张量上定位了具体位置，才发现是输入文本里有某个超长字符序列导致 embedding 的值分布被拉偏。

另一个很隐蔽的坑是 optimizer 的参数 epsilon。Adam 系列的 epsilon 默认值是 1e-8，这个在小模型上够用，但如果你用了混合精度，就必须调成 1e-6 甚至 1e-7，否则分母的数值稳定性会出问题，直接导致 loss 周期性变成 nan。我在普通 FP32 训练下没遇到这个问题，但这个点值得你记住，以后切混合精度时能少浪费一个下午。

4.3 显存与内存溢出的优化

从零手写模型比用封装好的 Trainer，更容易踩显存不足的坑，因为你对计算图的保留方式更“裸”。第一个优化方向是 padding。建议开启动态 batch，而不是固定 batch size。比如固定“每个 batch 的 token 总数”为 20000，短样本的 batch 可以容纳 64 条，长样本可能只有 8 条。这样显存利用率远比固定 batch 高，训练速度也更快。

第二个方向是梯度累积和torch.no_grad()评估。评估时如果你忘了包no_grad()，会有一部分显存被冻结在评估阶段的计算图里，几个 epoch 下来直接溢出。这个问题我用torch.cuda.max_memory_allocated()前后对比过，多占了接近 30% 显存。

第三个方向是避免过多中间变量。在多头注意力里，q、k、v三个矩阵的view和transpose会产生新的张量，如果中间结果不参与最终的梯度计算，尽量复用变量名，或者用torch.utils.checkpoint做重计算来换显存。我这个项目因为模型足够小，没有用到 checkpoint，但如果你后续放大模型，这个技术非常值得掌握。

4.4 复现性设置与微调细节

from scratch 项目最怕的就是同一个代码跑两回结果不一样。虽然用 PyTorch 时绝对值很难完全一致，至少要让实验结论可复现。我从这次项目中总结了一套相对可靠的复现配置清单：

这个表不是万能的，因为某些操作如torch.nn.functional.scaled_dot_product_attention在 GPU 上有浮点非确定性问题。但至少能保证你在同一台机器上两次训练的主要趋势一致，足以做超参数对比实验。

此外我还踩过一个小坑：DataLoader 的shuffle=True会在每次 epoch 重新打乱样本顺序，但如果你在训练中途恢复 checkpoint，训练样本顺序会变化，严格说数据集分布不完全一致。想彻底复现的话，需要自己在 sampler 里维护一个 index 数组，并把它也存进 checkpoint。

我个人在实际操作中最深的体会是：从零实现一遍不是为了证明你不需要框架，而是为了让你用框架时心里有底。以前我看训练日志里那些 loss 曲线，只是扫一眼看有没有下降；现在我看到的是每个模块在这条曲线背后到底做了什么，哪里可能出问题，哪里可以动刀。这种掌控感，是任何封装好的工具都给不了的。

这套代码跑通之后，我又往里加了几个“额外功能”——比如把每层的 attention score 分布导出出来观察衰减情况、把梯度范数曲线和 loss 曲线对齐可视化。这些调试手段会越来越成为你的直觉。如果你也打算照这个思路搭一套 from scratch 工程，建议你从一个小分类任务起步，别一上来就怼大模型。项目越小，你能踩的坑越完整。后面真想扩展，往模型层加个 encoder、往训练层加个混合精度、往推理层加个 batch 动态 padding，都是顺理成章的事。