前几篇笔记把自注意力拆了个底朝天Q、K、V 怎么乘多头怎么切softmax 的温度对分布有什么影响都算有了一些直观认识。但真到自己动手把一个叫 MiniMind 的小型语言模型往 GPU 上塞的时候才发现光盯着注意力那一亩三分地根本不够。训练跑不起来、推理吐字慢、长一点的上下文直接爆显存这些问题几乎都出在“注意力之外”。位置信息怎么塞进模型、前面的词怎么“记住”、显存和时间的账怎么算、注意力层之间的那些模块到底在干什么——它们不显眼但每一个都在决定模型能不能用、好不好用。这篇文章就把注意力之外这几块一次性讲透包含可直接用的显存计算公式、KV Cache 配置思路、以及我在调 MiniMind 时踩过的几个实在坑。适合正在复现小型语言模型、或者准备从原理走向实训的开发者哪怕你只是啃过注意力论文、还没动手写过训练循环也能从中拿到一些能直接上手的判断依据。1. 位置编码自注意力天生没有“先后”概念顺序到底怎么塞进去1.1 一个反直觉的实验打乱 token 顺序输出完全不变先做一个让人印象深刻的实验构造一个单层自注意力块输入“我 爱 学习”和“学习 爱 我”这两串 token不做任何位置处理直接过注意力层。你会发现两串输入产生的表示几乎完全相同在给定同样的输入嵌入后仅顺序不同输出中对应位置的向量仍然一致因为注意力对每个 token 的计算是加权求和顺序只影响“谁和谁加权”在图结构里这叫作置换不变性。这个特性在深度学习中不总是坏事但对语言模型来说就是致命伤语言的含义极度依赖词序。没有位置编码模型根本无法区分“A 打了 B”和“B 打了 A”。所以位置编码的本质不是锦上添花而是把离散的顺序信息转成模型能学习、能反向传播的连续向量。你可以把这个理解成给每个 token 发一个座位号。座位号不能太随意它要满足两个基本条件一是同一个位置在不同序列里语义一致二是位置之间要存在某个“距离感”让模型能感知到“这两个词隔了多远”。1.2 正弦位置编码与可学习位置编码经典方案的取舍早期的 Transformer 架构给了两种选择固定的正弦编码和可学习的嵌入表。正弦编码不用训练直接给位置 i 生成一个向量第 d 维用 sin 或 cos 的周期函数表示import math def sinusoidal_embeddings(seq_len, d_model): pe torch.zeros(seq_len, d_model) position torch.arange(0, seq_len, dtypetorch.float).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe.unsqueeze(0)为什么用不同频率的正弦原因是任意固定偏移 k 的位置向量都可以用其他位置的线性组合近似表达这给模型提供了一种“相对位置”的归纳偏置。模型不需要死记绝对位置而是可以通过线性变换感知两个 token 之间的相对距离。可学习位置编码则简单粗暴初始化一张[max_seq_len, hidden_dim]的表随训练一起更新。它的好处是灵活让模型自己去发现合适的位置表示坏处是外推差——一旦训练时没超过 2048遇到 4096 的序列就不知道怎么编码了。我个人的实践结论是在小模型参数在亿级以下上两种方案最终 loss 差距很小正弦编码的外推略好一些可学习编码在短序列任务上收敛稍快。如果你要处理长度变化很大的数据或者懒得出新版本先用正弦编码几乎不会错。1.3 旋转位置编码RoPE为什么现在主流这两年越来越多开源模型改用旋转位置编码RoPE原因很简单它在不增加参数的情况下把位置信息编码进了 Q 和 K 的内积里让注意力分数天然携带相对位置信息。RoPE 的思路可以这样类比想象你有一个二维向量每移动一个位置这个向量就旋转固定的角度。两个向量做内积时结果只取决于它们的相对旋转角差。扩展到高维就是对 Q 和 K 的向量按维度两两分组每组根据位置旋转不同角度。实际实现时并不需要真的把整个矩阵旋转而是用“对半拆分再点乘”的等价形式def rotate_half(x): x1, x2 x.chunk(2, dim-1) return torch.cat((-x2, x1), dim-1) def apply_rotary_pos_emb(q, k, cos, sin): return q * cos rotate_half(q) * sin, k * cos rotate_half(k) * sin我复现 MiniMind 时第一次手写 RoPE犯了一个典型错误只在 Q 上做了旋转忘记对 K 做同样处理结果注意力分数被强制引入了一个方向性的偏差训练 loss 一直不稳。排查了很久才发现是 K 忘了旋转。这个对称性非常关键RoPE 的核心是让 Q 和 K 的相对旋转角参与内积任何一边缺失都会破坏整个设计。为什么选择 RoPE 而不是其他方案核心优势是长度外推友好。由于相对位置靠“旋转角度”表达而旋转角度对位置差是线性的模型在更长序列上即使没见过也能给出还算合理的注意力分布。实测同一个 0.3B 小模型训练长度 1024用正弦编码在长度 2048 验证集上困惑度明显升高RoPE 只是小幅波动这个差距在实际生成任务里能直接感知到。2. 记忆KV Cache 与上下文窗口的账本2.1 自回归生成迫使我们缓存 K 和 V自回归语言模型的推理是一个 token 一个 token 蹦出来的。生成第 100 个 token 时理论上需要让第 100 个 token 和前面 99 个 token 做完整注意力计算。如果每次都从头开始算一遍计算量会随序列长度二次增长总计算量和序列长度的平方成正比生成 1024 个 token 就比 512 个慢约 4 倍。但仔细一看前 99 个 token 的 K 和 V 向量其实在新 token 出现之前就已经算好了而且它们不会因为新 token 的到来而改变自回归模型的因果掩码决定了这一点。所以一个标准优化是把已经算过的 K、V 存起来每次生成新 token 时只计算新 token 的 Q、K、V然后拿新 Q 去和缓存里所有的 K 做注意力。这个缓存就叫 KV Cache。这里要强调一点KV Cache 缓存的是每一层的 K 和 V不是某一块统一的大矩阵。每个 Transformer 层都有自己的 K/V 投影矩阵输出后都要缓存一份。层数越多缓存总量越大。很多框架里看到的past_key_values就是一个嵌套结构外层是层序号内层是 K 和 V。2.2 KV Cache 的显存公式算一笔实在账KV Cache 的显存体量非常直观每层 KV 缓存显存 2K 和 V 两份 × batch_size × seq_len × hidden_dim × 位宽 总显存 每层显存 × num_layers以我给 MiniMind 配的一个中等规模配置为例hidden_dim 1024层数 8推理时 batch 1序列长度 2048BF16 位宽占 2 字节单层 2 × 1 × 2048 × 1024 × 2 8,388,608 字节 ≈ 8MB 总共 8 × 8MB 64MB看起来不吓人对吧但把 batch 从 1 提到 32序列提到 8192加起来就是 3.4GB 以上2 × 32 × 8192 × 1024 × 2 × 8 8.5GB这还只是 KV Cache不算模型参数和激活。真正长上下文场景里KV Cache 往往是显存里涨得最快的部分。我调试时见过一个怪现象模型参数才 2GB跑 4000 上下文 batch 一上去直接 OOM一查全是 KV Cache 吃的。应对方案有几个维度减少头数GQA分组查询注意力通过让多个 Q 头共享一组 K/V 头把 KV 规模压缩到原来的 1/4 或 1/8训练不太影响质量推理显存骤降。滑动窗口只让 token 与最近 N 个 token 做注意力缓存也随之限制在窗口内。这个方案牺牲长程依赖但适合流式场景。量化缓存将 K/V 从 BF16 压到 INT8显存直接减半代价是精度损失需要评估实际任务。2.3 长上下文不是无代价的接触过“长上下文”能力的人容易误会只要把位置编码改成 RoPE把训练序列拉长模型就能记住更多信息。实际上位置编码解决的是“能不能感知远位置”但 KV Cache 决定的是“推理时放不放得下”而训练时的上下文窗口还受激活显存和优化器状态共同制约。三个因素必须同时考虑。MiniMind 当时做了个小实验同一份代码把训练序列从 512 拉到 2048训练速度下降了近四倍原因不是算子变慢了而是显存不够导致 batch size 不得不减半加上长序列自注意力的计算量平方增长。所以如果你想做长上下文模型要有一个清醒的预期长序列是“系统性”地贵不是某一个模块贵。3. 省显存训练时显存都去哪了以及怎么抠3.1 显存的四本账参数、梯度、优化器、激活训练时的显存去向比推理复杂得多。简单概括有四大块模型参数模型权重本身。梯度反向传播算出来的梯度大小和参数一致。优化器状态如果用 Adam每个参数要额外存一阶动量 m 和二阶动量 v而且通常以 FP32 保存占用可达参数的两倍乘二。激活值前向传播时每层产生的中间结果反向传播算梯度时必须用到所以不能随手丢掉。激活和 batch × 序列长度成正比在长序列场景下经常是最大的单块开销。假设一个小型模型参数 5 亿BF16 训练参数约 1GB梯度 1GBAdam 状态约 4GBFP32 下每个参数 8 字节已经 6GB 了。如果 batch 和序列再大一点激活轻松多出几个 GB。也就是说在一张 16GB 的卡上这个规模的模型其实非常紧张。一个小技巧用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()打点看训练循环里显存是在前向阶段涨得多还是在优化器 step 阶段涨得多。这比猜高效得多。我的经验是大多数小模型项目首先爆在激活值上而不是参数。3.2 混合精度省一半显存但小心溢出BF16 和 FP16 都能让模型参数和激活少用一半显存。但两者有个关键差别FP16 的动态范围小梯度很容易在反向传播时下溢成 0BF16 牺牲了尾数精度但动态范围和 FP32 相近所以现在大模型训练普遍偏爱 BF16。但 BF16 不是没有代价。它在更新参数时精度粗糙加上损失值刚好很大或很小时优化器容易跑偏。最稳妥的组合是“BF16 做前向后向 FP32 做优化器更新 梯度全部归一到固定区间”这已经是社区里很成熟的模式了。我刚开始用 BF16 时天真地以为只要把模型half()就行结果训练两轮 loss 奇高无比。后来排查才发现我在torch.amp.autocast外面手动把输入转成了 half反而绕过了 autocast 内部对某些算子的 FP32 保底策略。混合精度不要自己手动 cast尽量把计算放在autocast上下文里让它自己决定哪些算子用低精度、哪些必须保持 FP32。3.3 梯度检查点用计算换显存何时划算梯度检查点gradient checkpointing的思路是前向传播时不保存所有激活值只每隔几层保存一个“检查点”反向传播需要某个激活时再重新计算。这是一个典型的“用时间换显存”策略。开启后激活显存通常能降到原来的 1/3 到 1/4但训练时间会增加 20% 到 40%。什么情况下值得开一个非常实用的判断指标如果 batch size 因为显存不够被压到很小导致 GPU 利用率已经明显下降那么开梯度检查点把 batch 提上去往往净收益更划算。反之如果当前 batch 已经跑得挺满开梯度检查点只会白白浪费时间。我在 MiniMind 上把层数升到 12、序列长度 1024 时8GB 卡不开检查点 batch 只能到 2GPU 利用率很低。开了检查点后 batch 提到 8训练吞吐反而提升了大约 15%显存峰值还降了几百 MB。这种反直觉的优化一定要实测对比不要只看纸面。4. 省时间注意力之外的计算瓶颈与算子级优化4.1 Flash Attention 解决的不仅仅是“快”“Flash Attention 很快”是共识但它为什么快标准注意力在 PyTorch 里的实现是一连串独立的算子QK^T 算一次、除以 sqrt(d)、softmax、再乘 V。每一步都会把中间矩阵写回显存HBM下一个算子再读出来。而 HBM 的带宽有限这些反复读写比矩阵乘本身花的时间更多。Flash Attention 的核心是把计算和访存融合起来把 Q、K、V 切成小块按分块顺序在 SRAM片上高速缓存里做完整的注意力计算通过在线 softmax 的递推公式维护全局归一化结果避免把大中间矩阵写回 HBM。它获得的加速不是来自减少浮点运算次数而是来自大幅减少对显存的读写。这给我们一个重要启发在小模型训练里算子融合思路比“减少计算量”更有效。我自己实现了一个朴素的注意力运算量比融合版本没差多少但速度慢了近两倍瓶颈就在每个中间结果的显存写回。如果你的项目不用现成框架把 QK^T、softmax、乘 V 合进一个自定义 kernel 或使用现成的融合注意力是性价比最高的加速手段。4.2 矩阵乘法之外的 IO 瓶颈很多刚接触模型训练的人盯着 FLOPs 看觉得算力是唯一的瓶颈。但实践中小模型训练常常先撞上两个隐形的墙数据加载耗时和 GPU 利用率。MiniMind 刚开始训练时log 显示每个 step 的 GPU 计算时间只有 200ms但整体 step 周期要 600 多毫秒。用 profiler 一查数据加载和 tokenize 占了大量时间。改成提前把数据 token 化并缓存成二进制文件再用DataLoader的num_workers预取step 周期立刻落到 300ms 左右训练提速将近一倍。另外小模型的参数量不大但序列长度和 batch 决定的总 token 数才是吞吐的关键。想要提升模型训练速度一个非常立竿见影的招数是用“动态 padding”把长句和短句分桶到不同 batch而不是让所有样本都补到最长长度。一个 1024 序列的 batch 如果混了大量短句浪费的计算量相当可观。4.3 推理阶段和训练阶段各自的时间黑洞推理和训练的时间瓶颈不一样。训练阶段大部分时间是矩阵乘法和反向传播优化手段集中在融合算子和增大有效 batch。推理阶段则分为 prefill一次处理整段输入和 decode一个 token 一个 token 生成。prefill 像训练一样吃算力decode 阶段则受制于访存带宽每次生成一个 token 都要把所有参数从头读一遍计算量不大但数据搬运量大GPU 很多时候是在“等数据”而不是“算数据”。所以推理优化有两板斧一是让每次 decode 的序列更长——把多个请求拼成一个 batch让一次参数读取服务更多 token这个思路也叫连续批处理二是用 KV Cache 和算子融合减少额外访存。MiniMind 在单 batch 推理时几乎跑不满 GPU但把 8 个用户请求拼批后吞吐能翻三倍这就是访存型推理场景的典型特征。5. 深加工FFN、归一化和残差这些“配角”其实决定上限5.1 FFN 是模型记忆知识的地方注意力层负责“从上下文里取信息”但它本身几乎没有可存储知识的空间——它的参数只是投影矩阵处理的是 token 之间的关系。真正承担知识存储和特征变换的是每个 Transformer 块里的前馈网络FFN。FFN 的经典结构就是先把向量从 hidden_dim 升到 4 倍甚至更高经过非线性激活再降回来class FFN(nn.Module): def __init__(self, hidden_dim, intermediate_dim): super().__init__() self.w1 nn.Linear(hidden_dim, intermediate_dim) self.w2 nn.Linear(intermediate_dim, hidden_dim) def forward(self, x): return self.w2(activation(self.w1(x)))为什么先升维再降维一种通俗的理解是低维空间里线性不可分的特征先映射到高维通常更容易被分开完成非线性变换后再压回原维度。大部分参数集中在 FFN 里也因此常常有人说“FFN 就是模型的记忆体”。在 MiniMind 上做过一个对比把 FFN 的中间维度从 1024 加到 2048loss 下降明显训练速度只慢了不到 20%但把注意力头数翻倍loss 无明显改善速度却慢了不少。这说明在小模型里把“深加工”的资源多给 FFN往往比多堆注意力头更划算。5.2 从 ReLU 到 SwiGLU激活函数的工程演进早期 Transformer 的 FFN 激活函数是 ReLU简单、稀疏、好算。后来出现 GELU 这类平滑激活函数梯度更好传。现在很多新模型用的是 SwiGLU它把输入先通过两路线性变换再用一个“门控”机制决定放行多少信息SwiGLU(x) activation(xW1) * (xW3)注意这里比经典 FFN 多了一个独立的 W3 投影。这个多出来的投影不是花架子它相当于给 FFN 增加了一层门控让模型对“哪些信息值得进入高维空间”有了更细的可控性。代价是参数量和计算量略增。工程上需要留意SwiGLU 的 FFN 中间矩阵有 W1、W2、W3 三份权重文件和前向代码要对应好。我在往 MiniMind 里接一个开源权重的性格测试时因为把中间维度算错了层数导致权重 mismatch排查了半天才发现是激活函数结构不同导致网络宽度对不上。5.3 归一化与残差深层网络不崩的秘密深层网络训练不崩靠的是残差连接和归一化一起发力。残差连接等于给每一层的输入和输出搭了一条“短路”梯度可以沿着这条短路直接回流避免深层网络梯度消失。归一化则是把每一层的输入分布拉回一个稳定区间保证训练过程中数据不会因为逐层放大而失控。位置编码、KV Cache、显存优化这些大多是在解决“能不能跑起来”的问题而归一化和残差解决的是“深了之后还能不能训动”。从个人体会说如果 MiniMind 在 4 层到 8 层时遇到训练不收敛我会先检查残差连接的顺序和归一化位置而不是急着调学习率。主流实现里比较推荐的是 Pre-Norm先归一化再过注意力或 FFN最后加残差。它训练更稳定适合较大的模型。Post-Norm 是早期结构在层数不多时表达能力可能更好但到深层容易震荡。MiniMind 只有 8 层时Post-Norm 还能训扩到 12 层后开始出现 loss spikes切到 Pre-Norm 就稳了。如果你在复现一个小模型建议直接 Pre-Norm省一堆调参功夫。归一化层本身也有讲究。LayerNorm 会减去均值、除以标准差但额外的均值计算需要一次全局归约。RMSNorm 直接跳过减均值这一项用均方根做缩放速度和稳定性都不错很多新模型都采用了它。它的哲学是减去均值带来的平移不变性对很多任务是可有可无的省掉反而更省算力。5.4 参数初始化深度加工的地基深加工部分容易忽略的还有初始化。很多人直接用一个默认的 Linear 初始化到深层训练非常容易爆。合理的做法是把 FFN 里的输出投影 W2 初始化为很小的值甚至全 0保证网络初始状态接近恒等映射深层信息能顺畅流过。这个细节在复现任何小型语言模型时都值得先写进代码成本极低收益极大。我自己的 MiniMind 项目早期版本就是没注意初始化8 层网络前 100 步 loss 几乎不动还以为写了 bug。无意中把 FFN 输出投影的 scale 调小后loss 立刻开始正常下降。这个经验让我对“深加工”的每一个算子都保持敬畏代码里的默认初始化并不适合所有结构的随意叠加。最后再分享一个实际操作里的小技巧当你想验证位置编码、KV Cache、归一化这些模块有没有写对时最好的办法不是直接跑完整训练而是先在一个极小的随机数据上做单步前向和反向确认 loss 能稳定下降。然后再对比开启和关闭这个模块的 loss 曲线。MiniMind 里所有模块都是照这个顺序验证的它能帮你把“算法思路错了”和“代码实现错了”快速区分开。注意力之外的这些工程细节往往比注意力本身更考验一个开发者的耐心。