跳到主要内容

放大到千亿参数 — LLaMA 的三处改动、注意力的提速改造与专家分工

这一章讲三件事: 现代大模型对原版 Transformer 的三处「必改项」;注意力机制的省算力改造;以及混合专家怎么让 560 亿参数的模型只用 130 亿活跃参数干活。 前置: 上一章的注意力三步、层归一化、前馈层。这一章全是在它们上面动刀。

1. 为什么不能直接放大​

先交代一个事实:GPT-3 之后 OpenAI 不再开源模型,ChatGPT 和 GPT-4 的架构至今没有公开;学界真正能拆开研究的,是以 LLaMA 为代表的开源模型1。

LLaMA 与第 02 章原版结构的差异,书里归纳为三处:前置层归一化(并把归一化函数换成 RMSNorm)、激活函数(神经元的弯折开关)换成 SwiGLU、位置编码换成 RoPE——一种旋转式的位置嵌入(用旋转角度给位置编码)2。这三处不是审美偏好,各自对应原版结构在超大参数下的一个具体故障。逐个看。

改动一:RMSNorm。 原版层归一化要把数值调整到均值 0、方差(波动幅度)1,这需要先算均值再算方差。

RMSNorm 只用均方根(平方平均)重新缩放,少算一次均值。单看一层省得不多,但归一化在每层要跑两次,几十层乘上几千亿参数,这笔省略就是实打实的计算量——而效果与层归一化相当3。

改动二:SwiGLU。 原版前馈层用 ReLU 激活;SwiGLU 把输入拆成两路,一路过 Swish 函数、一路直接相乘。它的妙处在形状可调:参数 β 趋近 0 时 Swish 退化为线性函数,趋近无穷大时退化为 ReLU,取 1 时是光滑且非单调的函数——介于「平滑」和「分段」之间,在 PaLM 等模型上广泛应用,多数评测优于 ReLU4。

改动三:RoPE。 原版的位置编码加在输入层,模型深处对「距离」的感知要靠层层传递;RoPE 的出发点是用绝对位置编码的方式实现相对位置编码——对查询和键向量做旋转,旋转角度跟位置成正比。两个词的点积(对应位相乘再相加)只取决于它们的旋转角之差,「相距多远」就直接写进了注意力分数里;因为「旋转」这个名字而得名旋转式位置编码5。

它还有个工程红利:这种变换的数表很稀疏(大部分位置是零),可以用逐位相乘提高计算速度6。

三种规模的 LLaMA 各自的配置(层数、头数等训练前定好的设置),感受一下放大的量级7:

参数规模层数注意力头数嵌入宽度训练词元数
6.7B323240961.0 万亿
13.0B404051201.0 万亿
65.2B806481921.4 万亿

层数和头数同步翻倍、表示宽度也翻倍——「放大」是三个方向一起放,不是只加层数。

2. 注意力优化:四处省法​

第 02 章留了一个尾巴:自注意力(词看词的那步计算)的时间和存储复杂度与序列长度呈平方关系——序列翻倍,计算翻四倍。书里把优化分成四条路线8。

路线一:少算——稀疏注意力(只算值得算的那部分配对)。

对训练好的模型做分析发现,注意力那张打分表里很多格子是空的,于是可以限制 Query-Key 对的数量。基于位置的经典搭配有五种积木:全局、带状(只看邻居)、膨胀(隔空看)、随机、局部块;Star-Transformer、Longformer、BigBird 都是不同积木的组合。

BigBird 还给了一个理论旁证:稀疏注意力的编码器和解码器可以模拟任何图灵机——这在一定程度上解释了为什么砍掉大部分注意力对,效果依然不差9。

路线二:少搬——FlashAttention。 这条路不改数学,改的是数据搬运。加速卡的内存分级:全局内存——行话叫 HBM(一种大容量显存)——大而慢,H100 上有 80GB、读取速度 3.35TB/s;片上共享存储小而快,每个线程块只有 228KB10。标准注意力实现要把中间矩阵 S 和 P 写回 HBM 再读出来,这两个中间矩阵的尺寸随序列长度呈平方增长,计算还没崩,搬运先把速度耗光了。FlashAttention 的做法是把输入切块,在片上以增量方式算 Softmax,只存归一化因子不存中间矩阵;反向传播时用重算代替读取。代价是纯计算量增加,但省下的搬运时间远多于多算的开销——总体更快,显存更省11。

路线三:省缓存——多查询注意力(MQA)。 推理时要缓存所有历史词元的键和值(KV 缓存,第 12 章细算这笔账)。多头注意力里每个头都有自己的一套 K 和 V;MQA 让所有头共享一份键和值,只各自保留查询——键值矩阵从 N 份变 1 份,显存占用大幅下降。Falcon、StarCoder 等模型采用了它;后续研究还发现可以对训练好的模型微调补上 MQA,只需要约 5% 的原始训练数据量12。

路线四:压缓存——多头潜在注意力(MLA)。 DeepSeek-V2 引入。MQA 是「共享」,MLA 是「压缩」:把键和值联合压缩成一个低维潜在向量,缓存里只存压缩后的版本(压缩后远小于原始键值总量),用时再展开。书里引述的理论分析表明:在 KV 缓存大小相同的前提下,MLA 的表现力优于分组共享的 GQA;配套的 TransMLA 方法还能把现成的 GQA 模型转换成 MLA 模型13。

四条路线合起来看是一条清晰的账本逻辑:注意力贵在两处——算力(平方对数)和搬运/存储(中间矩阵与 KV 缓存)。稀疏注意力省算力,FlashAttention 省搬运,MQA 和 MLA 省缓存。

3. 混合专家:参数多不等于算得多​

到此为止有一条铁律:模型参数越多,每个词元过一遍的计算就越多。混合专家模型——行话叫 MoE——打破的正是这条绑定。依据是缩放法则说「参数是性能的关键」,但预算有限——于是让不同参数块(专家)专注不同任务或数据,每次只用一小部分14。

MoE 层由两部分组成:一个门控(打分选择)网络——也叫路由——负责决定把输入发给哪些专家,和 N 个专家网络。它通常替换掉 Transformer 块里的前馈层——选择这里不奇怪,在 5400 亿参数的 PaLM 里,90% 的参数都在前馈层15。

主走查:一个词元在 MoE 层的旅程​

拿 Mixtral-8x7B 来走。它有 8 个专家,门控网络每次激活 2 个;总参数量约 560 亿,但每个词元只用到约 130 亿活跃参数,性能却能比肩 700 亿参数的稠密(全部参数都参与干活的)模型 Llama-2-70B16。

  输入词元:「猫」(经注意力层后的向量)
      │
      ├─ 门控网络给 8 个专家打分(softmax 前):
      │    专家1: 3.1   专家2: 0.2   专家3: 2.8   专家4: -1.0
      │    专家5: 0.4   专家6: 1.9   专家7: -0.3  专家8: 0.1
      │
      ├─ TopK 保留前 2 名(专家1、专家3),其余设为 -∞
      │  (softmax 后 -∞ 变成 0,等于没参加)
      ├─ softmax → 专家1 权重 0.58,专家3 权重 0.42(演示数值)
      │
      └─ 输出 = 0.58 × 专家1(「猫」) + 0.42 × 专家3(「猫」)
                其余 6 个专家对这个词元一个数都没算。

  图说:TopK 和加权求和都是书里的机制;打分数值是为演示编的。
  省算力的机制一目了然:8 个专家只跑 2 个,参数躺在显存里不干活。

两个工程细节。其一,训练 MoE 时常给门控分数加噪声(随机扰动):一方面制造探索空间促使专家分工,

另一方面打破局部最优(卡在眼前的小坑里出不来)、提高训练稳定性17。其二,多个专家可能重复学到同样的通用知识,造成参数冗余;DeepSeekMoE 的解法是分出若干共享专家——不管路由怎么选,每个词元都先过共享专家,把「大家都要用的知识」收拢到一处,其余路由专家保持总数不变(激活数减去同等数量)18。

书里还给了另外两种变体,取舍很清楚:稠密 MoE 激活全部专家,不省计算(主要用于 LoRA+MoE 的微调场景,靠插件保住世界知识);软 MoE(SMEAR)干脆用门控权重把所有专家的参数加权平均成一个融合前馈层,只算这一个——计算成本接近单专家,但省掉了路由的不确定性。书里算了细账:软 MoE 的计算复杂度是 (L×4+N×2)×d×m,而稠密 MoE 是它的 N 倍19。

4. 作者的判断与证据​

  • 有证据的: LLaMA 的配置表(即超参数——训练前人工定好的设置——清单)、Mixtral 总参数/活跃参数比、PaLM 前馈层参数占比,书里都给了文献出处。
  • 作者的判断: 「绝大多数大模型采用 GPT 式 decoder-only 结构」是书里的现状归纳,不是定理——编码器-解码器结构在翻译类任务中仍有一席之地。
  • 书里没展开的: RMSNorm 为什么「少算均值不掉点」、SwiGLU 为什么有效,书里只给现象和出处,没有机理解释。这是这个领域的常态:很多组件是「试出来有效」,不是「推出来必须」。

5. 边界与局限​

  • MLA 和 GQA 的表现力对比依赖理论论文的假设条件,「MLA 更强」不是在所有任务上无条件成立。
  • MoE 的专家分工是涌现(无人设计、自己长出来)的、不可控的:书里说门控「依据数据特性精准路由」,但没有机制保证专家真的按人类理解的「领域」分工;这是开放问题(书里未讨论,我们的补充)。
  • 稀疏注意力的五积木组合是 BERT 时代(千级词元)的产物;当代长文本模型主要走 FlashAttention 加压缩路线,原书的稀疏注意力综述有时代局限。

6. 可带走的​

  1. 现代开源大模型 ≠ 原版 Transformer:RMSNorm(省一次均值)、SwiGLU(可调形状的激活)、RoPE(把相对距离写进注意力分数)是三处常见改造;
  2. RoPE 的本质:对查询和键做与位置成正比的旋转,点积自动只依赖位置差——这为第 06 章的位置插值埋下伏笔;
  3. 注意力的四条省钱路:少算(稀疏)、少搬(FlashAttention)、共享缓存(MQA)、压缩缓存(MLA);
  4. 算力和搬运是两本账:FlashAttention 计算量反而增加,但省下的内存搬运让它更快——优化前先搞清楚瓶颈在哪一侧;
  5. KV 缓存是推理时代最重要的显存消耗,MQA/MLA 都是在砍它——第 12 章会回到这条线;
  6. MoE 的门控+TopK 让「参数量」和「每词元计算量」解耦:Mixtral 560 亿参数只跑 130 亿;
  7. 共享专家防冗余、门控噪声防局部最优、软 MoE 免路由——MoE 的工程细节全在对抗「专家不干活或干重活」;
  8. BigBird 的图灵机等价性给了一个安心丸:注意力砍到很稀疏,表达能力理论上不塌。

7. 原文地图​

主题原书章原文位置
GPT-3 后不开源2 大语言模型基础text/02-ch02.txt:777(搜「不再开源」)
LLaMA 三处改动2 大语言模型基础text/02-ch02.txt:783(搜「前置层归一化」)
RMSNorm 与前置归一化2 大语言模型基础text/02-ch02.txt:822(搜「前置层归一化方法」)
SwiGLU 形状与 β2 大语言模型基础text/02-ch02.txt:867(搜「趋近于 0」)
RoPE 复数旋转2 大语言模型基础text/02-ch02.txt:875(搜「复数的思想」) · text/02-ch02.txt:885(搜「旋转」)
稀疏性逐位相乘2 大语言模型基础text/02-ch02.txt:914(搜「逐位相乘」)
LLaMA 超参数表2 大语言模型基础text/02-ch02.txt:973(搜「不同规模的 LLaMA 模型使用的超参数」)
平方复杂度2 大语言模型基础text/02-ch02.txt:1036(搜「平方的关系」)
稀疏注意力与积木2 大语言模型基础text/02-ch02.txt:1048(搜「稀疏」) · text/02-ch02.txt:1051(搜「全局注意力」)
BigBird 图灵机2 大语言模型基础text/02-ch02.txt:1065(搜「带状注意力和全局注意力」)
GPU 内存分级2 大语言模型基础text/02-ch02.txt:1123(搜「228KB」)
中间矩阵 S、P2 大语言模型基础text/02-ch02.txt:1144(搜「中间矩阵」)
FlashAttention 机制2 大语言模型基础text/02-ch02.txt:1135(搜「FlashAttention」) · text/02-ch02.txt:1149(搜「FLOPS 增加」)
MQA 共享键值2 大语言模型基础text/02-ch02.txt:1192(搜「多查询注意力」)
MLA 低秩压缩2 大语言模型基础text/02-ch02.txt:1274(搜「多头潜在注意力」)
MoE 动机与门控2 大语言模型基础text/02-ch02.txt:1324(搜「混合专家」) · text/02-ch02.txt:1331(搜「门控网络」)
PaLM 90% 参数在前馈层2 大语言模型基础text/02-ch02.txt:1337(搜「90% 的参数」)
Mixtral 走查数据2 大语言模型基础text/02-ch02.txt:1414(搜「560 亿」)
共享专家2 大语言模型基础text/02-ch02.txt:1432(搜「共享专家」)
软 MoE 复杂度2 大语言模型基础text/02-ch02.txt:1343(搜「软混合专家模型」)

Footnotes​

  1. 出处:「2 大语言模型基础」第 777 段(text/02-ch02.txt:777,搜「不再开源」)。 ↩

  2. 出处:「2 大语言模型基础」第 783 段(text/02-ch02.txt:783,搜「前置层归一化」)。 ↩

  3. 出处:「2 大语言模型基础」第 822 段(text/02-ch02.txt:822,搜「前置层归一化方法」)与式 (2.19)-(2.20)。RMSNorm 只计算均方根并缩放,不再减均值;「省得不多但每层都省」是我们的算账角度。 ↩

  4. 出处:「2 大语言模型基础」第 856 段(text/02-ch02.txt:856,搜「SwiGLU」)与第 867 段(text/02-ch02.txt:867,搜「趋近于 0」)。 ↩

  5. 出处:「2 大语言模型基础」第 875 段(text/02-ch02.txt:875,搜「复数的思想」)与第 885 段(text/02-ch02.txt:885,搜「旋转」)。 ↩

  6. 出处:「2 大语言模型基础」第 914 段(text/02-ch02.txt:914,搜「逐位相乘」)。 ↩

  7. 出处:「2 大语言模型基础」表 2.1(text/02-ch02.txt:973,搜「不同规模的 LLaMA 模型使用的超参数」)。 ↩

  8. 出处:「2 大语言模型基础」第 1036 段(text/02-ch02.txt:1036,搜「平方的关系」)。 ↩

  9. 出处:「2 大语言模型基础」第 1065 段(text/02-ch02.txt:1065,搜「带状注意力和全局注意力」);五种积木见第 1051-1060 段(text/02-ch02.txt:1051,搜「全局注意力」)。 ↩

  10. 出处:「2 大语言模型基础」第 1123 段(text/02-ch02.txt:1123,搜「228KB」)。原文:H100 全局内存 80GB/3.35TB/s,共享存储 228KB 且速度快得多。 ↩

  11. 出处:「2 大语言模型基础」第 1135 段(text/02-ch02.txt:1135,搜「FlashAttention」)与第 1149 段(text/02-ch02.txt:1149,搜「FLOPS 增加」)。 ↩

  12. 出处:「2 大语言模型基础」第 1192 段(text/02-ch02.txt:1192,搜「多查询注意力」)与第 1196 段(text/02-ch02.txt:1196,搜「5% 的原始训练数据量」)。 ↩

  13. 出处:「2 大语言模型基础」第 1274 段(text/02-ch02.txt:1274,搜「多头潜在注意力」)与第 1318 段(text/02-ch02.txt:1318,搜「TransMLA」)。 ↩

  14. 出处:「2 大语言模型基础」第 1324 段(text/02-ch02.txt:1324,搜「混合专家」)。 ↩

  15. 出处:「2 大语言模型基础」第 1337 段(text/02-ch02.txt:1337,搜「90% 的参数」)。 ↩

  16. 出处:「2 大语言模型基础」第 1414 段(text/02-ch02.txt:1414,搜「560 亿」)。走查中的门控打分和权重数值为演示编造。 ↩

  17. 出处:「2 大语言模型基础」第 1413 段(text/02-ch02.txt:1413,搜「噪声项」)。 ↩

  18. 出处:「2 大语言模型基础」第 1432 段(text/02-ch02.txt:1432,搜「共享专家」)。 ↩

  19. 出处:「2 大语言模型基础」第 1343 段(text/02-ch02.txt:1343,搜「软混合专家模型」)与结尾的复杂度分析段。 ↩