論文地址：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}^{?}{\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(?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^{?}{\text{k}}_j^C$，
$$\begin{array}{rl}{\text{q}}_t^{?}{\text{k}}_j^C& =(W^Q{\text{h}}_t{)}^{?}{W}^{UK}{\text{c}}_j^{KV}\\ & ={\text{h}}_t^{?}{W^Q}^{?}{W}^{UK}{\text{c}}_j^{KV}\\ & ={\text{h}}_t^{?}{W}^{QUK}{\text{c}}_j^{KV}\end{array}\text{\hspace{1em}?}$$
推理時可以將${W^Q}^{?}{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^{?}{\text{k}}_j}{\sqrt{d_h}}){\text{v}}_j\\ & =W^O\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\text{q}}_t^{?}{\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^{?}{\text{k}}_j}{\sqrt{d_h}}){\text{c}}_j^{KV}\\ & =W^{OUV}\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\text{q}}_t^{?}{\text{k}}_j}{\sqrt{d_h}}){\text{c}}_j^{KV}\end{array}\text{\hspace{1em}?}$$
$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 }(?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^{?}f({\text{k}}_j^C,j)& =(W^Q{\text{h}}_t{)}^{?}{R}_j{W}^{UK}{\text{c}}_j^{KV}\\ & ={\text{h}}_t^{?}{W^Q}^{?}{R}_j{W}^{UK}{\text{c}}_j^{KV}\end{array}\text{\hspace{1em}?}$$
由于$R_j$是未知敏感的，導致${W^Q}^{?}{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}^{?}{\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}(?)$表示應用RoPE的操作；$[?;?]$表示拼接操作。在推理時，解耦的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)}(?)$和${\text{FFN}}_i^{(r)}(?)$表示第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(?)$是指示函數；$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。

? 結果。