1. 这不是又一篇“Transformer原理扫盲”而是一次亲手搭积木的过程你点开这篇文章大概率是因为被“Attention is All You Need”那篇论文的公式吓退过或者在某次技术分享里听到“自注意力机制”“位置编码”“多头拼接”这些词时脑子像被塞进一团毛线——每个词都认识连起来却读不懂。更常见的情况是你已经用过Hugging Face的pipeline跑通了文本生成但当别人问“为什么加个LayerNorm就能稳定训练”你只能含糊说“好像是为了归一化”。这很正常。我带过二十多个实习生几乎所有人第一次看Transformer架构图时第一反应都是盯着那个带箭头的“QKV”框发呆而不是去想它到底在解决什么问题。这篇博文要做的不是再给你讲一遍“Transformer由Encoder-Decoder组成”这种教科书定义而是带你从零开始用不到300行纯PyTorch代码实现一个真正能跑起来、能调试、能观察中间状态的极简语言模型——它没有BERT那么大不追求SOTA指标但它每一层、每一个张量形状、每一次矩阵乘法都清清楚楚摆在你面前。它的名字就叫TinyLM一个只有2层Encoder、嵌入维度64、词表大小5000、最大序列长度128的“玩具模型”。但它足够真实它会分词、会计算注意力权重、会在训练时反向传播梯度、会在验证时输出困惑度Perplexity。更重要的是它所有模块的命名和结构都严格对应原始论文里的组件——你调试TinyLM.encoder.layers[0].self_attn时看到的就是论文Figure 2里那个带“Multi-Head Attention”的方块。这个项目适合三类人一是刚学完线性代数和链式法则想把抽象数学映射到具体代码的学生二是工作三年以上、日常调参但想搞懂底层为什么这么设计的工程师三是教学者需要一个可打断、可打印、可逐行注释的课堂演示模型。它不承诺让你明天就复现LLaMA但它能确保你合上电脑时脑子里不再只有“黑箱”两个字而是浮现出q k.T / sqrt(d_k)这行代码执行时GPU显存里实际流动的数据形状与数值范围。接下来的所有内容都基于这个目标展开让Transformer从幻灯片上的示意图变成你IDE里可以print()、breakpoint()、甚至手动修改权重后观察效果变化的活体对象。2. 整体设计思路为什么选择“极简但完整”的路径2.1 不做“最小可行模型”而做“最小可理解模型”很多人尝试简化Transformer时会走向两个极端一种是砍掉所有非核心模块比如直接用nn.Linear替代Embedding层用固定随机数代替位置编码最后得到一个连输入输出格式都不匹配的“骨架”另一种是保留全部结构但用超小参数如d_model8结果训练完全不收敛你无法判断是代码bug还是规模太小导致的数值不稳定。TinyLM的设计哲学是牺牲规模不牺牲结构完整性降低复杂度不降低概念保真度。它必须包含且仅包含原始Transformer Encoder中定义的5个核心组件Token Embedding、Positional Encoding、Multi-Head Self-Attention、Feed-Forward Network、Layer Normalization。少一个你就无法建立与论文的映射多一个比如Dropout或Pre-LN结构就会增加干扰变量。为什么嵌入维度定为64因为这是能在CPU上流畅运行、同时又不会因数值过小导致梯度消失的临界点。我实测过d_model32前向传播没问题但反向传播时torch.autograd.gradcheck会报gradcheck failed原因是浮点精度下梯度值趋近于零而d_model128在无GPU的笔记本上单步训练耗时超过8秒打断调试节奏。64是一个经验平衡点——它让q k.T的结果矩阵形状为(128, 128)你可以轻松用plt.imshow()可视化注意力热力图也让nn.Linear(64, 256)的权重矩阵尺寸为64×256打印出来刚好占满终端屏幕宽度方便你肉眼核对初始化分布。2.2 数据流设计强制显式张量形状声明绝大多数教程在讲解Transformer时会把x self.ln1(x self.attn(x))这样的代码一笔带过却不告诉你x在此刻的shape是什么self.attn(x)返回的又是几个tensor。TinyLM在每个模块的forward函数开头都插入了形如# x: [B, T, C]的注释并在关键节点添加assert校验。例如在Self-Attention模块中def forward(self, x): # x: [B, T, C] where Bbatch_size, Tseq_len, Cembed_dim assert x.dim() 3 and x.shape[-1] self.embed_dim B, T, C x.size() # Compute Q, K, V q self.w_q(x) # [B, T, C] - [B, T, C] k self.w_k(x) # [B, T, C] - [B, T, C] v self.w_v(x) # [B, T, C] - [B, T, C] # Reshape for multi-head: [B, T, C] - [B, T, H, C//H] - [B, H, T, C//H] q q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) v v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # Now q: [B, H, T, C//H], k: [B, H, T, C//H], v: [B, H, T, C//H]这段代码的价值不在于实现多头注意力而在于它强迫你面对一个事实张量形状的每一次变换都是有明确物理意义的。transpose(1, 2)不是为了炫技而是为了让q和k能在H维度上并行计算T×T的注意力矩阵view操作中的C // self.n_head不是随意除法而是确保每个头分配到等量的通道数。我在教学中发现学员卡壳90%的原因不是不懂矩阵乘法而是不知道某个.permute()操作后q的第0维到底是batch还是head。TinyLM把这种“形状契约”写死在代码里用assert在运行时兜底比任何文字描述都管用。2.3 训练策略放弃AdamW回归SGD with Momentum当前主流教程几乎清一色使用AdamW优化器理由是“收敛快、调参简单”。但这恰恰掩盖了Transformer最精妙的设计之一LayerNorm如何与优化器协同工作来稳定训练。AdamW自带的二阶矩估计即自适应学习率会平滑掉梯度尖峰让你误以为LayerNorm只是“锦上添花”。而TinyLM强制使用torch.optim.SGD(model.parameters(), lr3e-4, momentum0.9)原因很现实当你把学习率调到1e-3时不用LayerNorm的模型会在第3个step就lossinf而加上LayerNorm后它能稳定跑到第100步。这种“生死一线间”的对比比一百页公式推导都直观。更关键的是SGD的梯度更新公式p p - lr * g是线性的你可以精确计算每一步权重的变化量。比如当我把self.w_q.weight的初始标准差设为0.02遵循GPT-2初始化用SGD训练10步后w_q.weight.grad.mean().item()大约是-0.0015这意味着权重平均向负方向移动了3e-4 * 0.0015 ≈ 4.5e-7——这个量级恰好在FP32精度的安全范围内。如果换成AdamWexp_avg和exp_avg_sq的指数衰减会让这个计算变得不可追溯。TinyLM选择SGD不是因为它更先进而是因为它更“透明”让你看清LayerNorm如何把原本可能溢出的梯度压缩到可管理的区间。3. 核心模块拆解从词嵌入到注意力权重的全程追踪3.1 Token Embedding不只是查表而是维度对齐的第一道关卡初学者常误以为Embedding层就是个“高级字典”输入token ID输出对应向量。但在Transformer中它的作用远不止于此。TinyLM的Embedding层定义如下self.tok_emb nn.Embedding(vocab_size5000, embedding_dim64) self.pos_emb nn.Embedding(block_size128, embedding_dim64) self.drop nn.Dropout(0.1)注意三个细节第一vocab_size5000不是随便选的。我用collections.Counter统计了WikiText-2验证集的词频发现出现次数≥5的token共4987个向上取整到5000既覆盖99.2%的词汇又避免因词表过大导致nn.Embedding参数爆炸5000×6432万参数而BERT-base是3万×768≈2300万第二block_size128意味着模型最多处理128个token的序列这直接决定了后续所有张量的T维度第三nn.Dropout(0.1)放在Embedding后而非Attention后这是原始论文的设定目的是在信息注入阶段就引入噪声增强泛化性。最关键的实践技巧在于嵌入层的初始化。很多教程直接用默认的nn.init.normal_但TinyLM采用GPT-2的方案nn.init.normal_(self.tok_emb.weight, mean0.0, std0.02)。为什么是0.02因为0.02² × 64 ≈ 0.0256即嵌入向量的L2范数期望值约为0.16这个量级与后续Linear层的初始化std0.02保持一致避免某一层的输出值远大于其他层导致梯度爆炸。我在调试时曾把std改成0.1结果第一个batch的loss就飙到1e8torch.isnan(loss)返回True——这就是维度对齐失效的直接后果。提示在训练初期务必监控self.tok_emb.weight.norm().item()。正常情况下它应该在0.1~0.3之间缓慢增长。如果低于0.05说明初始化过小模型学不动如果高于0.5说明初始化过大需要重新设置std。3.2 Positional Encoding正弦波不是玄学而是傅里叶基的工程妥协“为什么用sin/cos而不用learnable position embedding”这个问题困扰了我整整两周。直到我把TinyLM的pos_emb换成可学习参数用相同数据训练100步发现其验证loss比正弦版高12%且注意力热力图呈现明显的“对角线偏好”模型只关注相邻token。这才明白正弦波的本质是用最少的参数编码最长的距离关系。其公式PE(pos, 2i) sin(pos / 10000^(2i/d_model))中10000这个常数不是拍脑袋来的——它是exp(9.21)的近似而9.21 ≈ ln(10000)确保当pos增大时不同频率的正弦波能覆盖从“相邻”到“超长距离”的全尺度。TinyLM的实现刻意避开了torch.nn.Embedding而是用纯函数生成def get_sinusoid_encoding_table(n_position, d_hid): Sinusoid position encoding table def get_position_angle_vec(position): return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] sinusoid_table np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) sinusoid_table[:, 0::2] np.sin(sinusoid_table[:, 0::2]) # dim 2i sinusoid_table[:, 1::2] np.cos(sinusoid_table[:, 1::2]) # dim 2i1 return torch.FloatTensor(sinusoid_table).unsqueeze(0) # [1, n_position, d_hid] self.pos_emb_table get_sinusoid_encoding_table(128, 64)这段代码的价值在于它让你看到pos_emb_table是一个固定的、非参数的张量。当你执行x self.tok_emb(idx) self.pos_emb_table[:, :idx.size(1), :]时加法操作要求tok_emb输出的[B, T, C]与pos_emb_table切片后的[1, T, C]能广播。这解释了为什么位置编码必须是[1, T, C]而非[T, C]前者支持batch维度广播后者会导致B维度丢失。我在第一次实现时漏掉了.unsqueeze(0)结果x的shape变成[B, T, C] [T, C]PyTorch自动广播成[B, T, C] [1, T, C]看似正确但梯度回传时pos_emb_table无法更新——因为它的requires_gradFalse。这个bug花了我3小时才定位教训是永远显式检查张量shape不要依赖广播的“巧合”。3.3 Multi-Head Self-Attention从单头到多头的升维真相这是整个模型最易误解的部分。很多人以为“多头”就是并行跑几个单头Attention然后拼接但TinyLM的实现揭示了一个关键事实多头的本质是把高维特征空间分解为多个低维子空间进行独立建模再融合结果。其核心代码如下# Project to Q, K, V q self.w_q(x) # [B, T, C] - [B, T, C] k self.w_k(x) # [B, T, C] - [B, T, C] v self.w_v(x) # [B, T, C] - [B, T, C] # Split into heads q q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # [B, H, T, C//H] k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # [B, H, T, C//H] v v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # [B, H, T, C//H] # Scaled dot-product attention attn (q k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) # [B, H, T, T] attn F.softmax(attn, dim-1) # [B, H, T, T] y attn v # [B, H, T, C//H] # Re-assemble all head outputs side by side y y.transpose(1, 2).contiguous().view(B, T, C) # [B, T, C]重点看y y.transpose(1, 2).contiguous().view(B, T, C)这一行。transpose(1, 2)把[B, H, T, C//H]变回[B, T, H, C//H]contiguous()确保内存连续否则view会报错view(B, T, C)则把H个头的输出在C维度上拼接。这里C必须等于H × (C//H)即总嵌入维度必须能被头数整除。TinyLM设n_head4所以C//H 64//4 16每个头处理16维特征。这意味着模型不是在64维空间里找全局最优而是在4个16维子空间里分别找局部最优再合并——这大幅降低了优化难度也增强了模型对不同关系模式的捕捉能力比如一个头专注语法另一个头专注语义。注意attn v的计算中attn是[B, H, T, T]的概率矩阵v是[B, H, T, C//H]的值向量。矩阵乘法[T, T] [T, C//H]的结果是[T, C//H]这正是注意力机制的物理意义用T个位置的权重对T个位置的值向量加权求和得到当前T位置的“上下文感知表示”。如果你打印attn[0, 0, 0, :]第一个样本、第一个头、第一个位置的注意力权重会看到它是一个和为1的[128]向量峰值通常在位置0附近——这证明模型确实学会了“关注自己”。3.4 Feed-Forward Network两层MLP为何是“非线性放大器”FFN层常被简化为“两个Linear加一个ReLU”但TinyLM的实现强调其作为非线性特征放大器的作用。其结构为self.mlp nn.Sequential( nn.Linear(embed_dim, 4 * embed_dim), # C - 4C nn.GELU(), nn.Linear(4 * embed_dim, embed_dim), # 4C - C nn.Dropout(0.1), )为什么隐藏层维度设为4 * embed_dim原始论文指出这是经验值旨在提供足够的容量来扩展特征表示。在TinyLM中64 → 256 → 64的映射让模型有能力将64维的“粗糙表示”先投影到256维的“精细特征空间”经过GELU激活比ReLU更平滑缓解梯度消失再压缩回64维。我做过对比实验把4 * embed_dim改成2 * embed_dim验证loss下降速度慢了40%改成8 * embed_dim训练时间翻倍但loss只改善2%性价比极低。GELU的选择也经过实测。用ReLU时约15%的神经元在训练初期就永久死亡输出恒为0而GELU的公式xΦ(x)Φ为标准正态CDF保证了所有输入都有非零梯度。你可以用torch.randn(10000) * 0.1生成输入计算F.relu(x).mean()和F.gelu(x).mean()前者接近0后者约为0.05——这微小的差异在深层网络中会被指数级放大。4. 实操全流程从环境搭建到loss曲线绘制的每一步4.1 环境准备与依赖安装拒绝“pip install transformers”TinyLM刻意避开Hugging Face生态所有依赖仅需torch2.0.1,numpy,matplotlib,tqdm。原因很实在当你想修改Attention的mask逻辑时不必在transformers源码的20层嵌套中找forward函数而是直接在自己的self_attn.py里改一行attn attn.masked_fill(mask 0, float(-inf))。安装命令极简pip install torch2.0.1 numpy matplotlib tqdm注意必须指定torch2.0.1。更高版本如2.1默认启用torch.compile会自动优化你的forward函数导致你无法用print()观察中间张量更低版本如1.13缺少torch.compile的fallback机制某些新API不可用。2.0.1是稳定性与功能性的最佳平衡点。4.2 数据预处理手写tokenizer比调用库更能理解分词本质TinyLM不使用tokenizers库而是用Python原生实现一个极简WordPiece tokenizerclass SimpleTokenizer: def __init__(self, vocab_size5000): self.vocab_size vocab_size self.stoi {} # string to int self.itos {} # int to string def train(self, text): # Step 1: split by whitespace and punctuation import re words re.findall(r\b\w\b, text.lower()) # Step 2: count frequency from collections import Counter counts Counter(words) # Step 3: take top vocab_size-3 (reserve 0PAD, 1UNK, 2EOS) most_common counts.most_common(vocab_size - 3) for idx, (word, _) in enumerate(most_common): self.stoi[word] idx 3 self.itos[idx 3] word self.stoi[PAD] 0 self.stoi[UNK] 1 self.stoi[EOS] 2 self.itos[0] PAD self.itos[1] UNK self.itos[2] EOS def encode(self, text): words re.findall(r\b\w\b, text.lower()) ids [self.stoi.get(w, 1) for w in words] # 1UNK return ids [2] # append EOS def decode(self, ids): return .join([self.itos.get(i, UNK) for i in ids])这个tokenizer的价值在于它让你直面分词的核心矛盾如何在有限词表下平衡OOVOut-of-Vocabulary率与序列长度。当你用tokenizer.train(wikitext_train)后执行len(tokenizer.stoi)会发现实际词表大小是4998预留2个给特殊token。此时若对验证集调用encode()统计UNK出现频率正常值应在0.8%~1.2%之间。如果超过2%说明词表太小需要增大vocab_size如果低于0.3%说明词表浪费可适当减小。我在调试时发现把re.findall(r\b\w\b)改成re.findall(r\w, ...)去掉\b边界符会导致“running”和“run”被当作不同词UNK率飙升至5%——这提醒你正则表达式的每个字符都在定义你的语言模型的认知边界。4.3 模型训练循环从零开始写trainer的硬核价值TinyLM的训练循环不调用Trainer而是手写train_step和val_stepdef train_step(model, data, targets, optimizer, criterion): model.train() optimizer.zero_grad() logits, loss model(data, targets) # logits: [B, T, vocab_size] loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() return loss.item() def val_step(model, data, targets, criterion): model.eval() with torch.no_grad(): logits, loss model(data, targets) return loss.item()关键点有三第一clip_grad_norm_的max_norm1.0不是随意设的。我测试过0.5和2.0max_norm0.5时梯度裁剪过于激进loss下降缓慢max_norm2.0时第50步开始出现lossnan。1.0是让95%的梯度范数落在[0.1, 0.9]区间的经验值。第二logits的shape是[B, T, vocab_size]而targets是[B, T]这要求criterion必须是nn.CrossEntropyLoss(ignore_index0)其中ignore_index0对应PAD避免padding token参与loss计算。第三model(data, targets)的实现中targets仅用于计算loss不参与前向传播——这是因果语言模型CLM的标准做法区别于掩码语言模型MLM。训练时我用tqdm包装epoch循环并实时打印for epoch in tqdm(range(100)): total_loss 0 for step, (data, targets) in enumerate(train_loader): loss train_step(model, data, targets, opt, crit) total_loss loss if step % 10 0: print(fEpoch {epoch} | Step {step} | Loss {loss:.4f}) avg_loss total_loss / len(train_loader) print(fEpoch {epoch} | Avg Train Loss {avg_loss:.4f})这个循环的价值在于它让你看到loss是如何一步步下降的。正常情况下前10个epochloss应从初始的8.5随机猜测降到4.250个epoch后降到2.8100个epoch后稳定在2.3±0.1。如果loss在5.0附近震荡说明学习率太大如果loss缓慢下降但始终高于3.0可能是数据预处理有误如未lowercase。4.4 可视化与调试用matplotlib画出你的第一个注意力热力图训练完成后最重要的验证不是看loss数字而是亲眼看到模型是否真的学会了关注相关token。TinyLM提供visualize_attention函数def visualize_attention(model, tokenizer, text, layer_idx0, head_idx0): model.eval() ids tokenizer.encode(text)[:128] # truncate to block_size x torch.tensor([ids], dtypetorch.long) # [1, T] with torch.no_grad(): # Forward pass, hook into attention weights attn_weights [] def hook_fn(module, input, output): attn_weights.append(output[1]) # output[1] is attention weights handle model.encoder.layers[layer_idx].self_attn.register_forward_hook(hook_fn) logits model(x) handle.remove() # Plot the first heads attention for first sample plt.figure(figsize(10, 8)) sns.heatmap(attn_weights[0][0, head_idx].cpu(), xticklabelsids, yticklabelsids, cmapviridis) plt.title(fLayer {layer_idx}, Head {head_idx} Attention) plt.show()运行visualize_attention(model, tokenizer, The cat sat on the mat)你会看到一个[7, 7]的热力图7个token。理想情况下对角线self-attention应是亮色且“cat”行应高亮“sat”和“mat”列“sat”行应高亮“cat”和“on”列。如果热力图全黑或全白说明模型未收敛如果只有对角线亮说明模型没学会长程依赖。我在调试时发现当pos_emb初始化为全零时热力图呈现完美的对角线——这证明位置信息缺失导致模型只能关注自身无法建立token间关系。5. 常见问题与排查技巧那些文档里不会写的坑5.1 “Loss is nan”浮点溢出的七种死法与解法这是TinyLM训练中最常遇到的问题90%源于以下原因问题根源表现特征解决方案原理说明Softmax输入过大attn矩阵中有inf或nan在softmax前加attn attn.masked_fill(mask 0, float(-inf))防止exp(large_number)溢出LayerNorm输入方差为0ln.weight梯度为nan初始化ln.weight为1.0ln.bias为0.0避免1/0除零错误Embedding层OOV过多loss初始值10检查tokenizer.encode()输出确保UNK率2%OOV token的logits无意义导致loss爆炸学习率过大loss在1e5量级震荡将lr从3e-4降至1e-4用lr_scheduler线性warmup大学习率使权重更新幅度过大Gradient explosiongrad.norm()100启用clip_grad_norm_(max_norm1.0)限制梯度范数防止权重突变Positional encoding未归一化pos_emb_table的norm()1.0用F.normalize(pos_emb_table, dim-1)防止位置编码主导token嵌入Batch size1时BN失效lossnan且batch_size1改用LayerNorm禁用所有BatchNormBN在batch1时方差为0导致除零我踩过的最深的坑是第七条在早期版本中我误用了nn.BatchNorm1d替代LayerNorm当batch_size1时BN计算var(x)得到0导致1/sqrt(var)为inf。解决方案不是换batch size而是彻底移除BN——Transformer的设计哲学就是不依赖batch统计量所有归一化必须是实例级的。5.2 “Attention weights are uniform”模型没学会关注的诊断树当你画出注意力热力图发现所有值都接近1/T如T128时约为0.0078说明模型处于“均匀注意力”状态尚未学到有效模式。按此顺序排查检查数据质量用print(tokenizer.decode(ids[:20]))查看前20个token确认是否为乱码或重复符号如全是PAD。如果是说明tokenizer.train()未正确执行。验证Embedding层执行print(model.tok_emb.weight[0, :5])应输出类似tensor([-0.012, 0.008, -0.021, 0.015, 0.003])的随机值。如果全为0说明初始化失败。检查QKV投影在self_attn.forward中插入print(q[0, 0, :5])应看到非零值。如果为0检查self.w_q权重是否被zero_grad()意外清零。确认Mask逻辑打印mask[0, :10, :10]应为下三角矩阵左上角10×10块。如果全1说明causal mask未生效模型能看到未来token破坏训练目标。验证Softmaxprint(attn[0, 0, 0, :5])应输出和为1的向量。如果和为0说明masked_fill填入了-inf而softmax对-inf返回nan。这个诊断树的价值在于它把抽象的“模型没学会”转化为具体的、可测量的张量值。我在修复一个bug时按此流程走到第4步发现mask是全1追查到torch.tril(torch.ones(T, T))被误写为torch.triu——一个字母之差让模型变成了“作弊者”。5.3 “Training is too slow”CPU/GPU瓶颈定位与加速技巧在无GPU环境下TinyLM单步训练耗时应≤1.2秒B16, T128, C64。如果超过2秒按此优先级优化首要检查torch.set_num_threads(6)。默认PyTorch使用所有CPU核心但线程竞争反而降低效率。设为物理核心数如i7-8700K是6核可提速35%。关闭梯度检查注释掉torch.autograd.gradcheck调用。它在每次backward时做数值梯度验证耗时是正常backward的10倍。减少日志打印将print(fStep {step} | Loss {loss:.4f})改为每10步打印一次。频繁I/O会阻塞计算。启用torch.compile在模型定义后添加model torch.compile(model)。它会将Python代码编译为优化的Triton内核在CPU上提速2.1倍实测。数据加载优化用DataLoader(..., num_workers2, pin_memoryTrue)。pin_memory让数据预加载到GPU显存减少CPU-GPU传输延迟。最有效的技巧是第五条。pin_memoryTrue本身不加速但它与torch.compile协同工作当数据已锁定在显存compile能更高效地调度计算图。我在一台16GB RAM的MacBook Pro上开启这两项后训练速度从1.8s/step降至0.7s/step提升157%。5.4 “Model outputs gibberish”生成文本质量差的根因分析当用model.generate(The cat)得到The cat a
 —— 用宝轻量框架设计必读)
