一文通透DeepSeek-V2(改造Transformer的中文模型)：从DeepSeek LLM到DeepSeek-V2的MLA与MoE

前言

成就本文有以下三个因素

24年5.17日，我在我司一课程「大模型与多模态论文100篇」里问道：大家希望我们还讲哪些论文
一学员朋友小栗说：幻方发布的deepseek-v2
24年5.24日，我司一课程「大模型项目开发线上营1」里的一学员朋友问我：校长最近开始搞deepseek了吗？刚看了论文，没搞懂MLA那块的cache是怎么算的，我总觉得他的效果应该类似MQA才对，但是反馈是挺好的

我当时回复他道：目前团队项目上的事情太多，然后近期在写那个KAN
确实还没来得及看这个deepseek，我近期看下
我们在继英文层面的论文翻译、审稿、对话、idea提炼之后
打算再整一下中文层面的「硕士论文修订助手(不限学科，CS和非CS的都涵盖)」
对于模型的选择，项目组有同事提议DeepSeek，故准备基于这个DeepSeek搞下

而搞之前——近几天，会先写一下它的论文解读，故本文就来了

且一如既往做到——对于几乎每一个主题，都如本博客万千读者或七月学员所说的：“还是看校长的文章好理解”

如有任何问题或任何不懂的地方，可以随时留言/评论，我会找时间尽快回复

第一部分

// 待更

第二部分 DeepSeek-V2之MLA

2.1 DeepSeek-V2提出MLA的背景与作用

DeepSeek-V2由国内量化公司深度求索发布，参数规模为236B，其中每个token激活21B的参数，且支持128K的上下文，其对应论文为《DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model》

为提高其性能

他们构建了一个高质量的多源预训练语料库，包括8.1T的token，与DeepSeek 67B使用的语料库相比，该语料库的数据量有所增加，特别是中文数据，并且数据质量更高
他们首先在完整的预训练语料库上预训练DeepSeek-V2
然后，收集了150万个对话会话，涵盖了数学、代码、写作、推理、安全等各个领域，以对DeepSeek-V2 Chat（SFT）进行监督微调（SFT）
最后，我们遵循DeepSeekMath的方法，采用组相对策略优化（GRPO）进一步使模型与人类偏好对齐，并生成DeepSeek-V2 Chat（RL）

DeepSeek-V2主要有两大创新点，其在transformer架构的基础上

通过创造性的提出Multi-head Latent Attent——简称MLA，改进了传统多头注意力(Multi Head Attention)，其本质意义便是降低了KV Cache的开销
且把FFN的结构改成DeepseekMoE——是对传统MoE结构的改进

2.1.1 KV Cache所导致的显存消耗大，需要尽可能降低

众所周知，KV Cache是大模型标配的推理加速功能——也是推理过程中，显存资源巨大开销的元凶之一。如下图所示，在模型推理时，KV Cache在显存占用量可达30%以上

目前大部分针对KV Cache的优化工作

比如著名的vLLM(这是其介绍页面、这是其对应的GitHub、其论文则为：Efficient Memory Management for Large Language Model Serving with PagedAttention，当然了，我后面会专门写一篇解读vLLM的博客)，其基于paged Attention，最大限度地利用碎片化显存空间，从而提升了空间利用率
再比如GQA、MQA
GQA是query数不变，但多个query组成一个group以共享一个key value
MQA则query也不变，但所有query共享一个key、一个value
至于更多，详见此文：一文通透各种注意力：从多头注意力MHA到分组查询注意力GQA、多查询注意力MQA

这类方案的问题是什么呢？在于比如前者并没有从根本上改变KV Cache占用空间巨大的问题，后者中的MQA虽然较大降低了KV cache计算量，但性能相比MHA下降太多了

那KV Cache到底是什么样的一个基本原理呢

对此，我们先来回顾下transformer当中的注意力计算公式(如对transformer还不够熟练，请看此文：Transformer通俗笔记：从Word2Vec、Seq2Seq逐步理解到GPT、BERT)
$\operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V$
GPT预测下一个token时，其只能看到待预测token之前的所有token，故在最终生成 $Q_1,Q_2,Q_3,Q_4$ 整个序列的过程中，会涉及到如下计算过程
$\operatorname{softmax}\left(\begin{array}{cccc} Q_{1} K_{1}^{T} & -\infty & -\infty & -\infty \\ Q_{2} K_{1}^{T} & Q_{2} K_{2}^{T} & -\infty & -\infty \\ Q_{3} K_{1}^{T} & Q_{3} K_{2}^{T} & Q_{3} K_{3}^{T} & -\infty \\ Q_{4} K_{1}^{T} & Q_{4} K_{2}^{T} & Q_{4} K_{3}^{T} & Q_{4} K_{4}^{T} \end{array}\right)\left[\begin{array}{c} \overrightarrow{V_{1}} \\ \overrightarrow{V_{2}} \\ \overrightarrow{V_{3}} \\ \overrightarrow{V_{4}} \end{array}\right]$
然后把上面的softmax结果和对应的V值一相乘，便可得到
$\begin{array}{l} \operatorname{Att}_{1}(Q, K, V)=\operatorname{softmaxed}\left(Q_{1} K_{1}^{T}\right) \overrightarrow{V_{1}} \\ \operatorname{Att}_{2}(Q, K, V)=\operatorname{softmaxed}\left(Q_{2} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmaxed}\left(Q_{2} K_{2}^{T}\right) \overrightarrow{V_{2}} \\ \operatorname{Att}_{3}(Q, K, V)=\operatorname{softmaxed}\left(Q_{3} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmaxed}\left(Q_{3} K_{2}^{T}\right) \overrightarrow{V_{2}} + \operatorname{softmaxed}\left(Q_{3} K_{3}^{T}\right) \overrightarrow{V_{3}} \\ \operatorname{Att}_{4}(Q, K, V)=\operatorname{softmaxed}\left(Q_{4} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmaxed}\left(Q_{4} K_{2}^{T}\right) \overrightarrow{V_{2}} + \operatorname{softmaxed}\left(Q_{3} K_{3}^{T}\right) \overrightarrow{V_{3}} + \operatorname{softmaxed}\left(Q_{4} K_{4}^{T}\right) \overrightarrow{V_{4}} \end{array}$
可以很明显的看到，上述计算过程中，有不少的 $K V$ 重复计算，比如 $K_{1}^{T}\overrightarrow{V_{1}}$ 、 $K_{2}^{T}\overrightarrow{V_{2}}$ 、 $K_{3}^{T}\overrightarrow{V_{3}}$

如果序列长度越长，类似这样的 $K V$ 重复计算会越多，从而势必将白白消耗那么大的显存，所以才说需要降低这种 $K V$ 重复计算

2.1.2 Multi-head Latent Attent：致力于在推理中降低 $n_{h} d_{h}$

MLA是对传统多头注意力做的改进，其目的有两个：首先是，降低推理过程中的KV Cache资源开销，其次，缓解MQA、MGA对性能的损耗

如上文所说，KV Cache中，提到每一步都需要将K和V缓存下来。假设单个Attention Block块中的多头注意力，有 $n$ 个头(或 $n_h$ )，每个 $k$ 和 $v$ 的维度为 $d$ (或 $d_h$ )，则每一步需要缓存的参数量为 $2 n_{h} d_{h} l$ ， $l$ 表示为transformer的层数
因此，MLA致力于在推理中降低 $n_{h} d_{h}$ ，具体而言，其不直接减少cache数量，而是类似Lora微调方法
对Key和Value进行了一个低秩联合压缩(通过低秩转换为一个压缩的KV，使得存储的KV的维度显著减小)，如下图所示

2.2 MLA的两个部分：一部分做压缩、一部分做RoPE编码

DeepSeek通过对Query和Key进行拆分为 $\left[q_{t}^{R}, q_{t}^{C}\right]$ 和 $\left[k_{t}^{R}, k_{t}^{C}\right]$ ，其中一部分做压缩 $\left(q_{t}^{C}, k_{t}^{C}\right)$ 、一部分做RoPE编码 $\left(q_{t}^{R}, k_{t}^{R}\right)$ (R可以理解为RoPE的标识符)

2.2.1 MLA对query和key的压缩

对于上图右下角的 $\mathbf{c}_{t}^{K V}$ 、 $\mathbf{k}_{t}^{C}$ 、 $\mathbf{v}_{t}^{C}$

可以看到先降维再升维

$\begin{array}{l} \mathbf{c}_{t}^{K V}=W^{D K V} \mathbf{h}_{t} \\ \mathbf{k}_{t}^{C}=W^{U K} \mathbf{c}_{t}^{K V} \\ \mathbf{v}_{t}^{C}=W^{U V} \mathbf{c}_{t}^{K V} \end{array}$

对于上述第一个公式， $c_{t}^{K V} \in R^{d_{c}}$ 是对key和value压缩后的隐向量，其通过一个降维映射矩阵 $W^{D K V} \in \mathbb{R}^{d_{c} \times d}$ 和模型输入 $h_t$ 「其中， $c_t$ 的维度 $d_c$ 远小于到头key和value的原始维度 $n_{h} d_{h}$ ，毕竟，别忘了上面说的，有 $n$ 个头，每个 $k$ 和 $v$ 的维度为 $d$ 」
至于上述第二三公式，在于通过第一个公式得到 $c_t$ 后，具体的key和value由两个对应的升维矩阵 $W^{U K}$ 和 $W^{U V}$ 还原

且在推理的过程中，只需要缓存每一步的 $\mathbf{c}_{t}^{K V}$ ，然后再计算还原回原始的K和V即可。由于 $c_t$ 的维度远小于K、V。因此每一步token的推理产生的缓存由之前的 $2 n_{h} d_{h} l$ ，变成 $d_c l$

对于上图左下角的 $c_{t}^{Q}$

之前提到KV Cache中，Q的作用只发生在当下(预测下一个token时，其只能看到待预测token之前的所有token)，但是在模型训练的过程中，每个输入的token会通过多头注意力机制生成对应的query、key和value

这些中间数据的维度往往非常高，因此占用的内存量也相应很大

所以论文中也提到，为了降低训练过程中的激活内存activation memory，DeepSeek-V2还对queries进行低秩压缩，即便这并不能降低KV Cache，而其对Q的压缩方式和K、V一致，依然是先降维再升维

$\begin{array}{l} c_{t}^{Q}=W^{D Q} h_{t} \\ q_{t}^{C}=W^{U Q} c_{t}^{Q} \end{array}$

其中

$\mathbf{c}_{t}^{Q} \in \mathbb{R}^{d_{c}^{\prime}}$ 是查询的压缩潜在向量(the compressed latent vector for queries)
$d_{c}^{\prime}\left(\ll d_{h} n_{h}\right)$ 表示查询压缩后的维度，而 $W^{D Q} \in \mathbb{R}^{d_{c}^{\prime} \times d}$ 、 $W^{U Q} \in \mathbb{R}^{d_{h} n_{h} \times d_{c}^{\prime}}$ 则分别表示查询的下投影和上投影矩阵(相当于先降维再升维)

2.2.2 MLA对query和key的RoPE编码

如下图红框所示，需要对 $\left(q_{t}^{R}, k_{t}^{R}\right)$ 做RoPE编码，并对其中的Key位置编码的部分进行Cache，从而在推理时不需要对Key进行位置编码的计算，提高了推理效率

在RoPE的实现中，如果要让Q、K带上位置信息，会分别乘以相应的位置编码矩阵

$\begin{array}{l} \tilde{Q}=R_{m} Q \\ \tilde{K}=R_{n} K \end{array}$

如果计算QK时，自然就变成了

$S=R_{m}^{T} Q^{T} R_{n} K$

通过上一节可知，DeepSeek-V2对Q和K都进行了压缩： $q_{t}^{C}=W^{U Q} c_{t}^{Q}$ 、 $\mathbf{k}_{t}^{C}=W^{U K} \mathbf{c}_{t}^{K V}$

分别做了如下压缩

$\begin{array}{l} \mathbf{c}_{t}^{K V}=W^{D K V} \mathbf{h}_{t} \\ \mathbf{k}_{t}^{C}=W^{U K} \mathbf{c}_{t}^{K V} \\ \mathbf{v}_{t}^{C}=W^{U V} \mathbf{c}_{t}^{K V} \end{array}$
$\begin{array}{l} c_{t}^{Q}=W^{D Q} h_{t} \\ q_{t}^{C}=W^{U Q} c_{t}^{Q} \end{array}$

则整个过程变成

$S=\left(W^{U Q}\right)^{T}\left(c_{t}^{Q}\right)^{T} R_{m}^{T} R_{n} c_{t}^{K V} W^{U K}$

其中的 $W^{U Q}$ 和 $W^{U K}$ ——如上节所述，分别是用于从低秩表示恢复到原始维度的解压缩矩阵(说白了，就是升维的)

可有个问题，问题在于，由于低秩表示已经是压缩了的状态，故直接在 $c_{t}^{Q}$ 和 $c_{t}^{K V}$ 上应用 $R_m$ 和 $R_n$ ，不等价于在完整的Q和K上应用位置编码(因为压缩操作可能已经丢失了某些信息，使得位置编码不能直接和有效地反映原始Q和K的位置关系)
为了解决这问题，Deepseek-V2设计了两个pe结尾的变量—— $\mathbf{q}_{t, i}^{R} \in \mathbb{R}^{d_{h}^{R}}$ 、 $\mathbf{k}_{t}^{R} \in \mathbb{R}^{d_{h}^{R}}$ (we propose the decoupled RoPE strategy that uses additional multi-head queries q𝑅𝑡,𝑖 ∈ R𝑑𝑅ℎ and a shared key k𝑅𝑡 ∈ R𝑑𝑅ℎ to carry RoPE, where $d_{h}^{R}$ denotes the per-head dimension of the decoupled queries and key，即表示解耦查询和键的每头维度)

用于储存旋转位置编码的信息，将信息存储和旋转编码解耦开
$\mathbf{q}_{t}^{R}=\operatorname{RoPE}\left(W^{Q R} \mathbf{c}_{t}^{Q}\right)$
$\mathbf{k}_{t}^{R}=\operatorname{RoPE}\left(W^{K R} \mathbf{h}_{t}\right)$
其中， $W^{Q R} \in \mathbb{R}^{d_{h}^{R} n_{h} \times d_{c}^{\prime}}$ 、 $W^{K R} \in \mathbb{R}^{d_{h}^{R} \times d}$ 分别是用于生成解耦查询和键的矩阵

压缩完、且RoPE编码完之后，最后将这4个变量—— $q_{t}^{C}=W^{U Q} c_{t}^{Q}$ 、 $\mathbf{k}_{t}^{C}=W^{U K} \mathbf{c}_{t}^{K V}$ 、 $\mathbf{q}_{t}^{R}$ 、 $\mathbf{k}_{t}^{R}$ ，分别拼接起来，形成带信息压缩的Q、K，以及带位置信息的Q、K，进行最后的计算

为方便大家更好的理解，上述这些公式可以对比下之前这个流程图

最终，单个Token产生的缓存包含了两个部分，即 $\left(d_{c}+d_{h}^{R}\right) l$

其中，如上文说过的， $n$ 个头(或 $n_h$ )，每个 $k$ 和 $v$ 的维度为 $d$ (或 $d_h$ )， $l$ 表示为transformer的层数， $n_g$ 表示为GQA中的组数， $d_{c}$ 、 $d_{h}^{R}$ 分别表示MLA中解耦查询和键的KV压缩维度和没被压缩的每头维度