论文地址：https://arxiv.org/pdf/2405.04434

相关博客
【自然语言处理】【大模型】DeepSeek-V2论文解析
【自然语言处理】【大模型】BitNet：用1-bit Transformer训练LLM
【自然语言处理】BitNet b1.58：1bit LLM时代
【自然语言处理】【长文本处理】RMT：能处理长度超过一百万token的Transformer
【自然语言处理】【大模型】MPT模型结构源码解析(单机版)
【自然语言处理】【大模型】ChatGLM-6B模型结构代码解析(单机版)
【自然语言处理】【大模型】BLOOM模型结构源码解析(单机版)

一、简介

• DeepSeek-V2是一个总参数为236B的MoE模型，每个token仅激活21B的参数，并支持128K的上下文长度。
• 提出了Multi-head Latent Attention(MLA)，通过压缩kv cache至隐向量，从而保证高效推理。
• 相比于DeepSeek 67B，DeepSeek-V2实现了更好的表现，节约了42.5%的训练成本，降低了93.3%的kv cache，提升最大吞吐5.76倍。
• 预训练语料包含了8.1T tokens并进一步进行SFT和RL。

二、模型结构

1. MLA(Multi-Head Latent Attention)

​ 传统Transformer采用MHA(Multi-Head Attention)，但是kv cache会成为推理瓶颈。MQA(Multi-Query Attention)和GQA(Grouped-Query Attention)可以一定程度减少kv cache，但效果上不如MHA。DeepSeek-V2设计了一种称为MLA(Multi-Head Latent Attention)的注意力机制。MLA通过低秩key-value联合压缩，实现了比MHA更好的效果并且需要的kv cache要小很多。

1.1 标准MHA

​ 令$d$为embedding维度，$n_h$是注意力头的数量，$d_h$是每个头的维度，${\text{h}}_t\in {\mathbb{R}}^d$是注意力层中第$t$个token的输入。标准MHA通过三个矩阵$W^Q,W^K,W^V\in {\mathbb{R}}^{d_h{n}_h\times d}$来产生${\text{q}}_t,{\text{k}}_t,{\text{v}}_t\in {\mathbb{R}}^{d_h{n}_h}$。
$$\begin{array}{rl}{\text{q}}_t& =W^Q{\text{h}}_t\\ {\text{k}}_t& =W^K{\text{h}}_t\\ {\text{v}}_t& =W^V{\text{h}}_t\end{array}\text{\hspace{1em}(1)(2)(3)}$$
在MHA中${\text{q}}_t,{\text{k}}_t,{\text{v}}_t$会被划分为$n_h$个头：
$$\begin{array}{rl}& [{\text{q}}_{t,1};{\text{q}}_{t,2};\dots ,{\text{q}}_{t,n_h}]={\text{q}}_t\\ & [{\text{k}}_{t,1};{\text{k}}_{t,2};\dots ,{\text{k}}_{t,n_h}]={\text{k}}_t\\ & [{\text{v}}_{t,1};{\text{v}}_{t,2};\dots ,{\text{v}}_{t,n_h}]={\text{v}}_t\\ & {\text{o}}_{t,i}=\sum _{j=1}^t\text{Softmax}(\frac{{\text{q}}_{t,i}^{\top }{\text{k}}_{j,i}}{\sqrt{d_h}}){\text{v}}_{j,i}\\ & {\text{u}}_t=W^O[{\text{o}}_{t,1};{\text{o}}_{t,2};\dots ,{\text{o}}_{t,n_h}]\end{array}\text{\hspace{1em}(4)(5)(6)(7)(8)}$$
其中${\text{q}}_{t,i},{\text{k}}_{t,i},{\text{v}}_{t,i}\in {\mathbb{R}}^{d_h}$是第$i$个注意力头的query、key和value，$W^O\in {\mathbb{R}}^{d\times d_h{n}_h}$是输出投影矩阵。在推理时，所有的key和value都会被缓存来加速推理。对于每个token，MHA需要缓存$2n_h{d}_{h}l$个元素。

1.2 低秩Key-Value联合压缩

​ MLA通过低秩联合压缩key和value来减少kv cache：
$$\begin{array}{rl}{\text{c}}_t^{KV}& =W^{DKV}{\text{h}}_t\\ {\text{k}}_t^C& =W^{UK}{\text{c}}_t^{KV}\\ {\text{v}}_t^C& =W^{UV}{\text{c}}_t^{KV}\end{array}\text{\hspace{1em}(9)(10)(11)}$$
其中${\text{c}}_t^{KV}\in {\mathbb{R}}^{d_c}$是用于压缩key和value的隐向量；$d_c(\ll d_h{n}_h)$表示KV压缩的维度；$W^{DKV}\in {\mathbb{R}}^{d_c\times d}$是下投影矩阵，$W^{UK},W^{UV}\in {\mathbb{R}}^{d_h{n}_h\times d_c}$表示上投影矩阵。在推理时，MLA仅需要缓存${\text{c}}_t^{KV}$，因此kv cache仅需要缓存$d_{c}l$个元素。此外，在推理时可以把$W^{UK}$吸收到$W^Q$，$W^{UV}$吸收到$W^O$中，这样甚至都不需要计算key和value。

关于推理时权重融合的理解：

​ 这里不考虑具体注意力头，仅就单个头进行分析。先来分析${\text{q}}_t^{\top }{\text{k}}_j^C$，
$$\begin{array}{rl}{\text{q}}_t^{\top }{\text{k}}_j^C& =(W^Q{\text{h}}_t{)}^{\top }{W}^{UK}{\text{c}}_j^{KV}\\ & ={\text{h}}_t^{\top }{W^Q}^{\top }{W}^{UK}{\text{c}}_j^{KV}\\ & ={\text{h}}_t^{\top }{W}^{QUK}{\text{c}}_j^{KV}\end{array}$$
推理时可以将${W^Q}^{\top }{W}^{UK}$预先计算出来，记为$W^{QUK}$。再来看整个注意力输出值的计算过程
$$\begin{array}{rl}{\text{u}}_t& =W^O{\text{o}}_t\\ & =W^O\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\text{q}}_t^{\top }{\text{k}}_j}{\sqrt{d_h}}){\text{v}}_j\\ & =W^O\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\text{q}}_t^{\top }{\text{k}}_j}{\sqrt{d_h}})W^{UV}{\text{c}}_j^{KV}\\ & =W^O{W}^{UV}\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\text{q}}_t^{\top }{\text{k}}_j}{\sqrt{d_h}}){\text{c}}_j^{KV}\\ & =W^{OUV}\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\text{q}}_t^{\top }{\text{k}}_j}{\sqrt{d_h}}){\text{c}}_j^{KV}\end{array}$$
$W^O{W}^{UV}$的结果记为$W^{OUV}$。通过这样的方式就不需要显式计算key和value。

​ 此外，为了在训练时降低激活的显存占用，对query也进行低秩压缩，即使其不能降低kv cache。具体来说，
$$\begin{array}{rl}{\text{c}}_t^Q& =W^{DQ}{\text{h}}_t\\ {\text{q}}_t^C& =W^{UQ}{\text{c}}_t^Q\end{array}\text{\hspace{1em}(12)(13)}$$
其中${\text{c}}_t^Q\in {\mathbb{R}}^{d_c^{\prime }}$是query的压缩后隐向量；$d_c^{\prime }(\ll d_h{n}_h)$表示query的压缩维度；$W^{DQ}\in {\mathbb{R}}^{d_c^{\prime }\times d},W^{UQ}\in {\mathbb{R}}^{d_h{n}_h\times d_c^{\prime }}$是下投影矩阵和上投影矩阵。

1.3 解耦RoPE

​ RoPE与低秩KV压缩并不兼容。具体来说，RoPE对于query和key是位置敏感的。若将RoPE应用在${\text{k}}_t^C$上，等式10中的$W^{UK}$将与位置敏感RoPE矩阵耦合。但是在推理时，$W^{UK}$就无法被吸收到$W^Q$中，因为对当前生成token相关的RoPE矩阵将位于$W^Q$和$W^{UK}$之间，而矩阵乘法不满足交换律。因此，推理时必须重新计算前面token的key，这会显著影响推理效率。

关于RoPE与抵秩KV压缩不兼容的理解。

RoPE向${\text{k}}_t^C$注入位置信息的方式为$f({\text{k}}_t^C,t)=R_t{\text{k}}_t^C$，其中$R_t$是一个位置敏感的矩阵。那么有
$$\begin{array}{rl}{\text{q}}_t^{\top }f({\text{k}}_j^C,j)& =(W^Q{\text{h}}_t{)}^{\top }{R}_j{W}^{UK}{\text{c}}_j^{KV}\\ & ={\text{h}}_t^{\top }{W^Q}^{\top }{R}_j{W}^{UK}{\text{c}}_j^{KV}\end{array}$$
由于$R_j$是未知敏感的，导致${W^Q}^{\top }{R}_j{W}^{UK}$针对不同的token，取值不一样。无法像先前那样直接融合为$W^{QUK}$。

​ 为了解决这个问题，提出使用额外的多头query ${\text{q}}_{t,i}^R\in {\mathbb{R}}^{d_h^R}$和共享key ${\text{k}}_t^R\in {\mathbb{R}}^{d_h^R}$来携带RoPE，其中$d_h^R$表示解耦query和key的每个头的维度。在MLA中使用解耦RoPE策略的方式为：
$$\begin{array}{rl}{\text{q}}_t^R& =[{\text{q}}_{t,1}^R;{\text{q}}_{t,2}^R;\dots ;{\text{q}}_{t,n_h}^R]=\text{RoPE}(W^{QR}{\text{c}}_t^Q)\\ {\text{k}}_t^R& =\text{RoPE}(W^{KR}{\text{h}}_t)\\ {\text{q}}_{t,i}& =[{\text{q}}_{t,i}^C;{\text{q}}_{t,i}^R]\\ {\text{k}}_{t,i}& =[{\text{k}}_{t,i}^C;{\text{k}}_t^R]\\ {\text{o}}_{t,i}& =\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\text{q}}_{t,i}^{\top }{\text{k}}_{j,i}}{\sqrt{d_h+d_h^R}}){\text{v}}_{j,i}^C\\ {\text{u}}_t& =W^O[{\text{o}}_{t,1};{\text{o}}_{t,2};\dots ;{\text{o}}_{t,n_h}]\end{array}\text{\hspace{1em}(14)(15)(16)(17)(18)(19)}$$
其中$W^{QR}\in {\mathbb{R}}^{d_h^R{n}_h\times d_c^{\prime }}$和$W^{KR}\in {\mathbb{R}}^{d_h^R\times d}$是用于产生解耦query和key的矩阵；$\text{RoPE}(\cdot )$表示应用RoPE的操作；$[\cdot ;\cdot ]$表示拼接操作。在推理时，解耦的key也需要被缓存。因此，DeekSeek-V2需要的总kv cache包含$(d_c+d_h^R)l$个元素。

1.4 结论

​ MLA能够通过更少的kv cache实现比MHA更好的效果。

2. 整体结构

2.1 基础结构

​ 对于FFN层，利用DeepSeekMoE架构，即将专家划分为更细粒度，从而获得更专业化的专家以及获取更准确的知识。在具有相同激活和总专家参数的情况下，DeepSeekMoE能够大幅度超越传统MoE架构。

​ 令${\text{u}}_t$是第t个token对FFN的输入，那么计算FFN的输出${\text{h}}_t^{\prime }$为：
h t ′ = u t + ∑ i = 1 N s FFN i ( s ) ( u t ) + ∑ i = 1 N r g i , t FFN i ( r ) ( u t ) g i , t = { s i , t , s i , t ∈ Topk ( { s j , t ∣ 1 ≤ j ≤ N r } , K r ) 0 , otherwise s i , t = Softmax i ( u t ⊤ e i ) \begin{align} \textbf{h}_t'&=\textbf{u}_t+\sum_{i=1}^{N_s}\text{FFN}_i^{(s)}(\textbf{u}_t)+\sum_{i=1}^{N_r}g_{i,t}\text{FFN}_{i}^{(r)}(\textbf{u}_t) \tag{20}\\ g_{i,t}&=\begin{cases} s_{i,t},& s_{i,t}\in\text{Topk}(\{s_{j,t}|1\leq j\leq N_r\},K_r)\\ 0,&\text{otherwise} \end{cases}\tag{21}\\ s_{i,t}&=\text{Softmax}_i(\textbf{u}_t^\top \textbf{e}_i) \tag{22}\\ \end{align} \\ ht′​gi,t​si,t​​=ut​+i=1∑Ns​​FFNi(s)​(ut​)+i=1∑Nr​​gi,t​FFNi(r)​(ut​)={si,t​,0,​si,t​∈Topk({sj,t​∣1≤j≤Nr​},Kr​)otherwise​=Softmaxi​(ut⊤​ei​)​(20)(21)(22)​
其中$N_s$和$N_r$表示共享专家和路由专家的数量；${\text{FFN}}_i^{(s)}(\cdot )$和${\text{FFN}}_i^{(r)}(\cdot )$表示第i个共享专家和第i个路由专家；$K_r$表示激活路由专家的数量；$g_{i,t}$是第i个专家的门限值；${\text{e}}_i$是当前层第i个路由专家的中心。

2.2 设备受限路由

​ 设计了一种设备受限路由机制来控制MoE相关的通信成本。当采用专家并行时，路由专家将分布在多个设备上。对于每个token，MoE相关的通信频率与目标专家覆盖的设备数量成正比。由于在DeepSeekMoE中细粒度专家划分，激活专家的数量会很大，因此应用专家并行时，与MoE相关的通信将更加昂贵。

​ 对于DeepSeek-V2，除了路由专家会选择top-K个以外，还会确保每个token的目标专家最多分布在M个设备上。具体来说，对于每个token，先选择包含最高分数专家的M个设备。然后在这M个设备上执行top-K选择。在实践中，当$M\ge 3$时，设备受限路由能够实现与不受限top-K路由大致一致的良好性能。

2.3 用于负载均衡的辅助loss

​ 不平衡的负载会增加路由坍缩的风险，使一些专家无法得到充分的训练和利用。此外，当使用专家并行时，不平衡的负载降低计算效率。在DeepSeek-V2训练时，设计了三种辅助损失函数用于控制专家级别负载均衡$({\mathcal{L}}_{\text{ExpBal}})$、设备级别负载均衡$({\mathcal{L}}_{\text{DevBal}})$和通信均衡${\mathcal{L}}_{\text{CommBal}}$。

​ 专家级均衡loss。专家级均衡loss用于缓解路由坍缩问题：
$$\begin{array}{rl}{\mathcal{L}}_{\text{ExpBal}}& ={\alpha }_1\sum _{i=1}^{N_r}{f}_i{P}_i,\\ f_i& =\frac{N_r}{K_{r}T}\sum _{t=1}^{T}1\text{(Token t selects Expert i)}\\ P_i& =\frac{1}{T}\sum _{t=1}^T{s}_{i,t}\end{array}\text{\hspace{1em}(23)(24)(25)}$$
其中${\alpha }_1$是称为专家级均衡因子的超参数；$1(\cdot )$是指示函数；$T$是序列中token的数量。

​ 设备级均衡loss。除了专家级均衡loss以外，也设计了设备级别均衡loss来确保跨设备均衡计算。在DeepSeek-V2训练过程中，将所有的专家划分至$D$组 { E 1 , E 2 , … , E D } \{\mathcal{E}_1,\mathcal{E}_2,\dots,\mathcal{E}_D\} {E1​,E2​,…,ED​}并在单个设备上部署每个组。设备级均衡loss计算如下：
$$\begin{array}{rl}{\mathcal{L}}_{\text{DevBal}}& ={\alpha }_2\sum _{i=1}^D{f}_i^{\prime }{P}_i^{\prime }\\ f_i^{\prime }& =\frac{1}{{\mathcal{E}}_i}\sum _{j\in {\mathcal{E}}_i}{f}_j\\ P_i^{\prime }& =\sum _{j\in {\mathcal{E}}_i}{P}_j\end{array}\text{\hspace{1em}(26)(27)(28)}$$
其中${\alpha }_2$是称为设备级均衡因子的超参数。

​ 通信均衡loss。通信均衡loss能够确保每个设备通信的均衡。虽然设备限制路由机制能够确保每个设备发送信息有上限，但是当某个设备比其他设备接收更多的tokens，那么实际通信效率将会有影响。为了缓解这个问题，设计了一种通信均衡loss如下：
$$\begin{array}{rl}{\mathcal{L}}_{\text{CommBal}}& ={\alpha }_3\sum _{t=1}^D{f}_i^{\prime \prime }{P}_i^{\prime \prime }\\ f_i^{\prime \prime }& =\frac{D}{MT}\sum _{t=1}^{T}1\text{(Token t is sent to Device i)}\\ P_i^{\prime \prime }& =\sum _{j\in {\mathcal{E}}_i}{P}_j\end{array}\text{\hspace{1em}(29)(30)(31)}$$
其中${\alpha }_3$是称为通信均衡因子的超参数。设备受限路由机制操作主要确保每个设备至多向其他设备传输MT个hidden states。同时，通信均衡loss用来鼓励每个设备从其他设备接受MT个hidden states。通信均衡loss确保设备间信息均衡交换，实现高效通信。

2.4 Token-Dropping策略

​ 虽然均衡loss的目标是确保均衡负载，但是其并不能严格确保负载均衡。为了进一步缓解由于不均衡导致的计算浪费，在训练时引入了设备级别的token-dropping策略。该方法会先计算每个设备的平均计算预算，这意味着每个设备的容量因子等于1.0。然而，在每个设备上drop具有最低affinity分数的token，直到达到计算预算。此外，确保大约10%的训练序列的token永远不会被drop。这样，可以根据效率要求灵活地决定是否在推理过程中drop token，并确保训练和推理的一致性。

三、预训练

1. 实验设置

1.1 数据构造

​ 数据处理过程同DeepSeek 67B，并进一步扩展数据量和质量。采用与DeepSeek 67B相同的tokenizer。预训练语料包含8.1T tokens，中文token比英文多12%。

1.2 超参数

​ 略

1.3 Infrastructures

​ DeepSeek-V2训练基于HAI-LLM框架。利用16路0气泡流水并行、8路专家并行和ZeRO-1数据并行。考虑到DeepSeek-V2具有相对较少的激活参数，并且对一部分操作进行重计算来节约激活显存，因此可以不使用张量并行，从而降低通信开销。此外，为了进一步提高训练效率，使用专家并行all-to-all通信来重叠共享专家的计算。使用定制化的CUDA核来改善通信、路由算法和不同专家之间融合线性计算。此外，MLA基于改善版本的FlashAttention-2进行优化。

1.4 长上下文扩展

​ 使用YaRN将上下文窗口尺寸从4K扩展至128K。

2. 评估

四、对齐

​ SFT。 使用了150万样本的微调数据，其中120万是用于有用性，30万则用于安全性。

​ 强化学习。仍然采用GRPO。

​ 结果。