DeepSeek Paper Note

涉及的论文

DeepSeek LLM: 并行训练

这些并行训练的东西目前已经变成了各家大模型的基础，这里先留空，稍后再更新，o7

DeepSeek V2: MLA 和 MoE 的高效推理

DeepSeek V2 的技术报告当中，有两大亮点。首先是 MLA 这个 Attention 变种，极大减小了 KV Cache，提高推理速度，并允许更长的上下文。另外是 MoE，将模型整体规模扩大的情况下，激活参数反而比 V1 更小，能够同时提高推理速度和质量。

MLA

科学空间的这篇文章，已经详细介绍了 MLA 的来龙去脉，并且有详细的公式推导。我们首先跟随论文的脚步，从 KV Cache 的角度，对论文中 -93.3% 的优化比例进行一下推导，以加深理解。首先，我们归纳一下 MLA 的特性，作为本节的 Takeaway。

MLA 的特性

KV Cache
1. MLA 不单独保存 Key（$k_i$）和 Value（$v_i$），而是保存更低维度的、由隐层向量$x_i$算出的$c_t$。
2. 减小 KV Cache 大小对推理速度是有利的，因推理过程是 Memory-Bounded，减小 KV Cache 大小可以减少数据传输延时。
3. KV Cache 需要“还原”为 Key 和 Value，有额外的计算步骤。
4. KV Cache 大小与 Head 数量无关，增加 Head 提升模型能力时，并不必然带来 KV Cache 的增加，于是 MLA 在 Scalability 方面有优势。
RoPE 需要单独处理：
1. $c_t$不包含位置编码，因此 concat 额外的$d_r$个维度，用于位置信息的计算。
2. 这些维度构成了 PE Cache，与 KV Cache 独立。
从 LoRA 的角度理解 MLA
1. MLA 可以理解成对 MHA 的 $W_{QKV}$ 进行了低秩分解，而 GQA 可以理解成 MLA 的一种特殊情况。
2. MLA 的 $c_t$ 相当于 LoRA 的中间层表示。
代码实现
1. 官方 Repo当中的实现符合上述特性。
2. HuggingFace目前的实现，会把 MLA 退化为一般的 MHA，以便调用 Flash Attention，不能享受 MLA 对 KV Cache 的优化。

KV Cache 的优化比例

首先，我们整理一下参与对比的两个模型的一些主要参数，以便计算 KV Cache 的大小。

参数	DeepSeek LLM 67B	DeepSeek V2 236B
KVCache 数据类型	BF16（16 bits）	量化后约 6 bits
层数 $l$	95	60
隐层维度 $d$	8192	5120

Attention 类型	GQA	MLA
Attention Head 数量 $h$	64	128
Attention Key 维度 $d_k$	128	128
Attention Value 维度 $d_v$	128	128
[GQA] Attention 组 $g$	8	-
[MLA] KV 潜在维度 $d_{ckv}$	-	512
[MLA] Q 潜在维度 $d_{cq}$	-	1536
[MLA] 位置编码维度 $d_r$	-	64

然后，计算元素数量
- V1：$(d_k + d_v)gl$ = 256 * 8 * 95 = 194560
- V2：$(d_{ckv} + d_r)l$ = 576 * 60 = 34560
- 比率：17.76%
最后，计算每个 Token 所需的 KV Cache 大小
- V1：194560 * 16 bits = 380 KiB
- V2：34560 * 6 bits $\approx$ 25.31 KiB
- 比率：6.66% （即 -93.34%）

MLA 的计算量分析

假设 KV Cache 目前的长度是 $m$ （上文已有$m$个 token），新增的 token 隐层表示为$x_t$。我们以乘法次数代表计算量，对三种 Attention 的计算量进行比较。这里参考了DeepSeek-V3的代码。

首先，我们大致算一下 MHA 的计算量，假设$d_k = d_v = d_q$且$d_k * h = d$
- 新 token 的 QKV：$(d * d_k + d * d_v + d * d_q) * h = 3dd_kh$
- Attention score 和 V 加权和：$(d_k * m) * h + (d_v * m) * h = 2d_khm$ （此处未计入 Softmax、scaling、Norm、 RoPE 的乘法次数）
- 回隐层表示维度：$d * d_v * h = dd_kh$
- 总计算量：$4dd_kh + 2d_khm = 4d^2 + 2dm$
然后，将它改成 GQA 的计算量，假设$d_k = d_v = d_q$且$d_k * h = d$
- 新 token 的 QKV：$d * d_k * g + d * d_v * g + d * d_q * h = 2dd_kg + dd_kh$
- Attention score 和 V 加权和：$(d_k * m) * h + (d_v * m) * h = 2d_khm$ （此处未计入 Softmax、scaling、Norm、 RoPE 的乘法次数）
- 回隐层表示维度：$d * d_v * h = dd_kh$
- 总计算量：$2dd_kh + 2dd_kg + 2d_khm = 2.25d^2 + 2dm$
最后，从头推导 MLA 的计算量，假设$d_k = d_v = d_q$，注意 MLA 并无$d_k * h = d$的约束
- 新 token 的 Q：$d * d_{cq} + d_{cq} * (d_q + d_r) * h = dd_{cq} + d_{cq}d_kh + d_{cq}d_rh$
- 新 token 的 $c_t^{kv}$：$d * (d_{ckv} + d_r) = dd_{ckv} + dd_r$
- KV Cache 还原 K，吸收到新 token 的 Q 中：$d_q * d_{ckv} * h = d_{ckv}d_kh$
- Attention score：$d_{ckv} * m * h + d_r * m * h = d_{ckv}hm + d_rhm$（此处未计入 Softmax、scaling、Norm、 RoPE 的乘法次数）
- KV Cache 还原 V，吸收到 V 加权和：$d_{ckv} * m * h + d_{ckv} * d_v * h = d_{ckv}hm + d_{ckv}d_kh$
- 回隐层表示维度：$d * d_v * h = dd_kh$
- 总计算量：$dd_{cq} + d_{cq}d_kh + d_{cq}d_rh + dd_{ckv} + dd_r + 2d_{ckv}d_kh + 2d_{ckv}hm + d_rhm + dd_kh$
- 按照次数分组：$(dd_{cq} + dd_{ckv} + dd_r) + (d_{cq}d_k + d_{cq}d_r + 2d_{ckv}d_k + dd_k)h + 17*(2d_{ckv} + d_r)hm$
这个值无法和前文的 MHA、GQA 的计算量直接比较，因为 MLA 的$h$是 GQA 的三倍之多，所以我们假设一个特殊的 MHA 和 GQA，$d_k = d * 1/40, h = 128$, 取$g = 16$以保持$h / g$不变
- 代入$d_k = d * 1/40，d_{cq}=d * 3/10，d_{ckv}=d * 1/10，d_r=d * 1/80, h=128$
- MHA: $4dd_kh + 2d_khm = 12.8d^2 + 6.4 dm$
- GQA: $2dd_kh + 2dd_kg + 2d_khm = 7.2d^2 + 6.4 dm$
- MLA：$33/80 d^2 + 33/800 * d^2h + 17 / 80dhm = 5.6925 d^2 + 27.2 dm$
- 这样，我们就对 MLA 的计算量有一个大致的认识：每个 Token 的计算量因分解而有所减少，但多出了将 KV Cache 还原的计算量。

MLA 的参数量分析

在 MHA 中，我们忽略 pre-norm 部分，则参数是四个矩阵，假设$d_k = d_v = d_q$且$d_k * h = d$
- $W_Q, W_K, W_V, W_O$：$4dd_kh$
然后，将它改成 GQA 的参数量
- $W_Q, W_O$：$2dd_kh$
- $W_K, W_V$：$2dd_kg$
- 总参数量：$2dd_kh + 2dd_kg$
最后，从头推导 MLA 的参数量，以官方代码中的符号为准，假设$d_k = d_v = d_q$，注意 MLA 并无$d_k * h = d$的约束
- $W_Q^A$：$dd_{cq}$
- $W_Q^B$：$d_{cq} * (d_q + d_r) * h = d_{cq}d_kh + d_{cq}d_rh$
- $W_{KV}^A$：$dd_{ckv} + dd_r$
- $W_{KV}^B$：$d_{ckv} * (d_k + d_v) * h = 2d_{ckv}d_kh$
- $W_O$: $d * d_v * h = dd_kh$
- $W_{KV}^A$和$W_{KV}^B$之间，$W_Q^A$和$W_Q^B$之间各有一个 RMS Norm 模块：$d_{ckv} + d_{cq}$，这部分参数量很小，因此忽略
- 总参数量：$dd_{cq} + d_{cq}d_kh + d_{cq}d_rh + dd_{ckv} + dd_r + 2d_{ckv}d_kh + dd_kh = (dd_{cq} + dd_{ckv} + dd_r) + (d_{cq}d_k + d_{cq}d_r + 2d_{ckv}d_k + dd_k)h$
同计算量的条件，进行代换
- 代入$d_k = d * 1/40，d_{cq}=d * 3/10，d_{ckv}=d * 1/10，d_r=d * 1/80, h=128, g = 16$
- MHA: $12.8 d^2$
- GQA: $6.4d^2 + 0.8d^2 = 7.2d^2$
- MLA: $33/80d^2 + 33/800d^2h = 5.6925d^2$
- 这样，我们就对 MLA 的参数量有一个大致的认识：和 GQA 的参数量大致相当，DeepSeek V2 的参数配置下略小于 GQA。

GLU

在 DeepSeek V2 中，每个 Expert 的结构是 GLU: Gated Linear Unit，在代码当中沿用了 MLP 的命名，但结构是不同的。如果是为了和以前的工作保持一致，可以将其称为 FFN。

GLU 的结构

从 llama 开始兴起的 GLU，相比于 $\text{MLP}(x) = W_{down}(\text{ReLU}(W_{up}(x)))$，GLU 在参数上多出一个门控矩阵，具体的分析见论文。它将输入数据分成两部分，一部分通过线性变换，另一部分通过门控函数进行激活，然后将这两部分相乘得到最终的输出。GLU 的公式如下：

$$ \text{GLU}(x) = W_{down}(\sigma(W_{gate}(x)) \odot W_{up}(x)) $$

$x$ 是输入数据
$W_{gate},W_{up}\in \mathbb{R}^{d\times d_h}$，$W_{down}\in \mathbb{R}^{d_h\times d}$，是线性变换矩阵权重矩阵
$\sigma$ 是 Sigmoid 函数
$\odot$ 是逐元素乘法

GLU 的参数量分析

对于每个 GLU 单元，其参数量是：$3dd_h$。

GLU 的计算量分析

对于每个 GLU 单元，其计算量是：$3dd_h + d_h$，按乘法次数计，激活函数的计算量未计入。

MoE

DeepSeek 使用的 MoE 网络，最早在DeepSeek MoE中提出。在这个知乎文章中，对三代 MoE 的区别做了详细的整理和解释。

MoE 的结构

MoE 是 Mixture of Experts 的缩写，即混合专家网络。MoE 的结构包括 2 个主要模块：

Expert：专家网络，每个专家网络可以独立地处理输入数据，并输出一个结果。
Gate：门控网络，负责选择哪些专家网络应该被激活，以及如何分配输入数据给这些专家网络。

这部分的公式如下，最终输出$y$是所有专家输出的加权和：

$$ \begin{align*} h' &= u + \sum_{i=1}^{N_r} g_{i} \text{FFN}_i(u), \\ g_{i} &= \frac{g'_{i}}{\sum_{j=1}^{N_r} g'_{j}}, \\ g'_{i} &= \begin{cases} s_{i}, & s_{i} \in \text{Topk}(\{s_{j}|1 \leq j \leq N\}, K), \\ 0, & \text{otherwise}, \end{cases} \\ s &= f(u^T e), \end{align*} $$

其中：

$u$是输入。
$N$是专家网络的数量。
$\text{FFN}_i(u)$是第$i$个专家网络的输出。
$f(u^T e)$是门控网络的输出，$f$是激活函数，一般是 SoftMax。
$g_i(x)$是门控网络对第$i$个专家的最终得分，经过了 TopK 选择和归一化。

MoE 是为了提升模型扩展性：

为了拟合更大的训练集合，存储更多的知识，需要更多的参数
为了模型的实用，需要减少激活的参数，从而减少计算量
每一个 Token 的预测，并不需要使用所有的知识，只需要很小的一部分参数就足以完成推理

这里的必然结果就是参数总量继续提升，但是每次推理只激活其中一部分参数，这就是“稀疏激活”的方法，而 MoE 是稀疏激活的一种实现。

DeepSeek MoE 的演化

DeepSeek MoE

在上文讲解的 MoE 基础结构上，DeepSeek MoE 做了 3 点改变：

【+】Fine-Grained Expert Segmentation：将 Expert 分的更小，激活的 Expert 数量更多，但保持总参数量和激活比例不变。
【+】Shared Expert Isolation：部分 Expert 始终激活。
【+】Load Balance：通过辅助 Loss 训练，使得 Expert 的激活尽可能均衡。包括：
1. Expert-Level：意图让每个 Expert 被激活的 token 数量均衡。
2. Device-Level：将 Expert 分组，每个组被激活的 token 数量均衡，每个组部署到同一个设备。

DeepSeek V2

在 DeepSeek MoE 上，DeepSeek V2 做了 3 点改变：

【+】Device-Limited Routing：每个 token 激活的 Expert，至多位于 M 个设备，设备的打分是其上的 Expert 打分之和。
【+】Communication Balance Loss：增加一个辅助 Loss，使得设备之间传输的 token 数量均衡，或者说，每个设备从其他设备收到的 token 数量均衡。
【+】Token-Dropping Strategy：为了进一步均衡训练负载，丢弃和当前设备的 Expert 的亲和性最低的一部分 token，这强制了训练阶段每个设备的负载均衡。另外，文章也提到，至少 10%的训练 token 始终保留。

DeepSeek V3

在 DeepSeek MoE 上，DeepSeek V2 做了 5 点改变：

【-】删除之前提出的所有辅助 Loss
【-】删除之前提出的 Token-Dropping
【~】Router 激活函数从 SoftMax 改为 Sigmoid
【+】Auxiliary-Loss-Free Load Balancing：给 Router 添加可学习的参数$b_i$。
1. 训练时：负载太高的 Expert，稍降低它的$b_i$。负载太低的 Expert，稍提高它的$b_i$。变动大小是超参数$\gamma$
2. $b_i$仅用于 TopK 选择，在后续的归一化和加权和中，不参与计算。
【+】Complementary Sequence-Wise Auxiliary Loss：意图让每个句子中的 token，均衡地激活各个 Expert。

总结

一口气看完三代 MoE 的变化，马后炮地来说：

负载均衡是 MoE 最重要的问题，不仅仅关乎计算效率，也关乎对参数的充分利用，最终会影响模型的能力。
Auxiliary Loss 会引入太多的超参数，是一个复杂的方法，而且可能和主 Loss 的优化方向冲突，应该不是负载均衡的最好方案。
Shared Expert Isolation 很漂亮，符合直觉。将一部分常用的知识提取到单独的 Expert，应该也有助于其他 Expert 的负载均衡。

MoE 的参数量分析

假设每个 Expert 的参数量为$M=3dd_h$，总共有$N$个 Expert，激活$K$个 Expert（Shared Expert 含在其中），那么：

Expert 参数量：$NM=3Ndd_h$
Router 参数量：$Nd$ （或者 $Nd+d$，若添加 bias）
总参数量：$NM+Nd = 3Ndd_h + Nd\approx 3Ndd_h$，主要参数量来自 Expert

MoE 的计算量分析

假设每个 Expert 的计算量为$M=3dd_h + d_h$，总共有$N$个 Expert，激活$K$个 Expert（Shared Expert 含在其中），那么：

Expert 计算量：$KM=3Kdd_h + Kd_h$
Router 计算量：$Nd$
总计算量：$KM+Nd=3Kdd_h + Kd_h + Nd\approx 3Kdd_h$，主要计算量来自被激活的 K 个 Expert

一般来讲，$K«N$，MoE 比同样参数量的 Dense 模型，计算量要少很多。