Deepseek-V2技術報告解讀！全網最細！

作者：Junda 2024-05-09 16:21:46

得力于出色的研究人員和工程團隊，Deepseek-V2將大語言模型訓練中廣泛被驗證有用的訓練策略深度整合，集合了長度外推訓練的YaRN，高效對齊的GRPO，MLA與混合專家分配等方法進行模型訓練。做到了算法、工程和數據的極致優化。

深度求索Deepseek近日發布了v2版本的模型，沿襲了1月發布的 Deepseek-MoE（混合專家模型）的技術路線，采用大量的小參數專家進行建模，同時在訓練和推理上加入了更多的優化。沿襲了一貫的作風，Deepseek對模型（基座和對話對齊版本）進行了完全的mit協議開源，可以商用。對于算力不是那么充足的開發者，官方提供了API調用的方案，費用更是達到了全場最低！

圖片

在技術報告的開始，Deepseek團隊用多個數字和兩張圖直觀地概括了目前模型取得的效果。模型參數量方面達到236B ，同時由于模型小專家混合的特性，模型在推理時的激活參數很少，可以實現高推理速度。在通用能力的表現上，模型在MMLU多選題benchmark上拿到78.5 分，取得了第二名，Deepseek-V2在眾多開源模型中表現僅次于70B 的 LLaMA3，超過了他們此前發布的V1代67B的非MoE模型。在成本效率方面，相比V1的稠密模型，V2模型節約了42.5%的訓練成本，減少了推理時93.3%的 KV-cache 顯存占用，將生成的吞吐量也提升到了原來的5.76倍。借助YaRN優化的長度外推訓練方法，模型的上下文能力得以擴展到了128k大小。下面我們結合代碼和技術報告，對Deepseek-V2模型進行詳細的解讀。

圖片

核心優化解析

在這里我們結合官方技術報告中的模型架構圖輔助說明，介紹模型的核心優化點——多頭隱式注意力（Multi-head Latent Attention,MLA）：

圖片

如上圖右下所示，大模型使用kv-cache進行模型的解碼加速，但是當序列較長的情況下很容易出現顯存不足的問題，MLA從這一角度出發，致力于減少kv緩存的占用。

圖片

MLA從LoRA的成功借鑒經驗，實現了比GQA這種通過復制參數壓縮矩陣尺度的方法更為節省的低秩推理，同時對模型的效果損耗不大。我們首先結合配置文件中的這幾行了解下每個部分的作用：

"hidden_size": 5120,
"kv_lora_rank": 512,
"moe_intermediate_size": 1536, 
"q_lora_rank": 1536,
"qk_nope_head_dim": 128,
"qk_rope_head_dim": 64

模型處理上一層計算出的隱藏狀態（hidden_size=5120）時，首先會將模型的q壓縮到 q_lora_rank這一維度（設定為1536），再擴展到 q_b_proj 的輸出維度（num_heads * q_head_dim），最后切分成 q_pe 和 q_nope 兩個部分，在訓練部分中我們將看到這樣設計的作用。

##### __init__ #####
self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim # =192
self.q_a_proj = nn.Linear(
   self.hidden_size, config.q_lora_rank, bias=config.attention_bias
)
self.q_a_layernorm = DeepseekV2RMSNorm(config.q_lora_rank)
self.q_b_proj = nn.Linear(
   config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
)
##### forward #####
bsz, q_len, _ = hidden_states.size()

q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
# q (bsz, q_len, 24576)
q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
# q (bsz, q_len, 128, 192)

q_nope, q_pe = torch.split(
    q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
)
# 將最后一層 192 的hidden_states切分為 128 (qk_nope_head_dim) + 64 (qk_rope_head_dim)

對于kv矩陣的設計，模型使用了kv壓縮矩陣設計（只有576維），在訓練時進行先降維再升維。在模型推理的時候，需要緩存的量變成 compressed_kv，經過 kv_b_proj 升高維度得到 k,v 的計算結果。

##### __init__ #####
self.kv_a_proj_with_mqa = nn.Linear(
  self.hidden_size,
  config.kv_lora_rank + config.qk_rope_head_dim,
  bias=config.attention_bias,
)
self.kv_a_layernorm = DeepseekV2RMSNorm(config.kv_lora_rank)
self.kv_b_proj = nn.Linear(
  config.kv_lora_rank,
  self.num_heads
  * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
  bias=False,
)
        
##### forward #####
compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
compressed_kv, k_pe = torch.split(
  compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
)
k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
kv = (
  self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
  .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
  .transpose(1, 2)
)

那么，為什么Deepseek-V2要把整個計算流程拆成 q_nope, k_nope, k_pe, k_nope 這四個部分呢？在RoPE的實現中，如果想要直接讓模型的 q, k 具有位置性質，通常是這樣做的，m,n 代表特定位置的token，R的含義可以查閱RoPE：

計算輸出的attention得分時，整個過程變成了：

為了節約KV cache的內存，Deepseek-V2將kv cache壓縮到了同一個小矩陣中，后面再解壓縮出來：

這個時候注意力得分的計算可以寫成：

這個時候我們變得清楚了，我們apply旋轉位置編碼的時候，標準的不帶解壓縮的實現是會將原始的K狀態直接更新到拼到K前面的，而上面的矩陣運算是使用先左乘，后解壓縮的方式，由于矩陣乘法是沒有交換律的，因此這種矩陣壓縮設定下使用C作為cache直接拼接在數學上是不等價的。為了解決這個問題，Deepseek-V2設計了兩個pe結尾的變量用于儲存旋轉位置編碼的信息，將信息存儲和旋轉編碼解耦合開。

之后，將q,k中負責儲存信息的部分，負責旋轉編碼的部分拼接起來，進行標準的attention計算：

k_nope, value_states = torch.split(
  kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
)
kv_seq_len = value_states.shape[-2]

cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
        
query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
query_states[:, :, :, self.qk_nope_head_dim :] = q_pe

key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
key_states[:, :, :, self.qk_nope_head_dim :] = k_pe
if past_key_value is not None:
    cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
    key_states, value_states = past_key_value.update(
        key_states, value_states, self.layer_idx, cache_kwargs
    )

attn_weights = (
  torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale
)

attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim)
attn_output = self.o_proj(attn_output)

最后將 num_head 維度拉平，經過輸出矩陣得到模型這一層的輸出隱藏狀態，仍為 5120 維。

架構解讀

我們通過模型的架構圖和配置文件對模型設計有一個大致的認知，Deepseek的模型習慣采用 remote_code導入的格式，下載模型后，我們通過官方示例導入模型權重，打印出模型的架構。

DeepseekForCausalLM(
  (model): DeepseekModel(
    (embed_tokens): Embedding(102400, 5120)
    (layers): ModuleList(
      (0): DeepseekDecoderLayer(
        (self_attn): DeepseekAttention(
          (q_a_proj): Linear(in_features=5120, out_features=1536, bias=False)
          (q_a_layernorm): DeepseekRMSNorm()
          (q_b_proj): Linear(in_features=1536, out_features=24576, bias=False)
          (kv_a_proj_with_mqa): Linear(in_features=5120, out_features=576, bias=False)
          (kv_a_layernorm): DeepseekRMSNorm()
          (kv_b_proj): Linear(in_features=5120, out_features=32768, bias=False)
          (o_proj): Linear(in_features=163840, out_features=5120, bias=False)
          (rotary_emb): DeepseekYarnRotaryEmbedding()
        )
        (mlp): DeepseekMLP(
          (gate_proj): Linear(in_features=5120, out_features=12288, bias=False)
          (up_proj): Linear(in_features=5120, out_features=12288, bias=False)
          (down_proj): Linear(in_features=12288, out_features=5120, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): DeepseekRMSNorm()
        (post_attention_layernorm): DeepseekRMSNorm()
      )
      (1-59): 59 x DeepseekDecoderLayer(
        (self_attn): DeepseekAttention(
          (q_a_proj): Linear(in_features=5120, out_features=1536, bias=False)
          (q_a_layernorm): DeepseekRMSNorm()
          (q_b_proj): Linear(in_features=1536, out_features=24576, bias=False)
          (kv_a_proj_with_mqa): Linear(in_features=5120, out_features=576, bias=False)
          (kv_a_layernorm): DeepseekRMSNorm()
          (kv_b_proj): Linear(in_features=5120, out_features=32768, bias=False)
          (o_proj): Linear(in_features=163840, out_features=5120, bias=False)
          (rotary_emb): DeepseekYarnRotaryEmbedding()
        )
        (mlp): DeepseekMoE(
          (experts): ModuleList(
            (0-159): 160 x DeepseekMLP(
              (gate_proj): Linear(in_features=5120, out_features=1536, bias=False)
              (up_proj): Linear(in_features=5120, out_features=1536, bias=False)
              (down_proj): Linear(in_features=1536, out_features=5120, bias=False)
              (act_fn): SiLU()
            )
          )
          (gate): MoEGate()
          (shared_experts): DeepseekMLP(
            (gate_proj): Linear(in_features=5120, out_features=3072, bias=False)
            (up_proj): Linear(in_features=5120, out_features=3072, bias=False)
            (down_proj): Linear(in_features=3072, out_features=5120, bias=False)
            (act_fn): SiLU()
          )
        )
        (input_layernorm): DeepseekRMSNorm()
        (post_attention_layernorm): DeepseekRMSNorm()
      )
    )
    (norm): DeepseekRMSNorm()
  )
  (lm_head): Linear(in_features=5120, out_features=102400, bias=False)
)

我們從上往下，從embedding層的維度來看，與Gemma, LLaMA和Qwen的經驗一致，Deepseek也選取了較大的輸入詞表作為模型的輸入（數據充足且多樣的情況下當然可以這么干），這樣做的好處是詞表的多樣性強，解碼的一個token內有多個字，壓縮效率很高。

"num_hidden_layers": 60,
"num_key_value_heads": 128,
"num_experts_per_tok": 6,
"n_shared_experts": 2,
"n_routed_experts": 160

通過以上配置分析，模型共有60個層，注意力頭數為128，總的門控專家個數為160，每個token計算有6個門控專家被激活，同時還有2個共享專家保持激活狀態，共計8個被激活的專家。在經過embedding層后，與Deepseek-MoE保持一致，首先會經過一個共享的大Decoder層進行第一層計算，這層模型的attention計算設定與后續59層基本一致，唯一區別是這一層的mlp層固定為8個專家的寬度，沒有門控額外參數激活的設定，這一設置與每層共享專家的設定一樣，研究者希望語言生成的公共知識（包含流暢性、邏輯性等）被存儲在這里。

而當我們從模型的整體架構選取上來看，層數足夠深的時候使用pre-norm方便模型訓練，歸一化使用RMSNorm，非線性激活函數使用SiLU，attention矩陣不加bias（對flash-attention有好處），這些似乎是如今大廠在訓練大模型時候會采用的標配了。

訓練

MLA設定下的解耦長度外推：模型使用基于進制轉換的YaRN進行長度外推訓練，在大海撈針測試中表現不錯。

模型對齊訓練：模型使用對話數據進行SFT，同時評估時重點關注指令遵循能力。在強化對齊階段也下了很大的功夫，最早出現在Deepseek-Math中的GRPO算法被用來進行偏好對齊訓練，這是一種無需在訓練中更新通常與Policy Model（被對齊模型）同樣大小的 Critic Model 的參數的訓練方法，是一種資源優化的 PPO。（注意：還是需要訓 Reward Model 的，只是不會在對齊的時候進行參數更新）

GRPO和PPO的對比

Infra

模型訓練的工程優化方面（infra）仍有很多給人啟示的點。模型使用了pp=16的流水線并行（pipeline parallel），160個專家分ep=8個節點并行（expert parallel），而并未采取任何形式的張量并行（tensor parallel），降低了通信成本，使用了ZeRO-1的數據并行來減少優化器狀態的顯存占用。訓練設施在卡間使用NVLink和NVSwitch，節點間使用InfiniBand交換機，通信優化已經全部拉滿。并行策略全部使用自研的HAI-LLM實現。

另外，Deepseek-V2結合算法和工程，提出了資源感知專家負載均衡的方法，保證了專家并行的幾個機器雨露均沾，不會出現有些機器空轉，有些機器過度占用的情況。在訓練時，結合模型本身的專家ensemble特性，各個專家在訓練開始的過程中是完全對稱的，這種設計如果不做額外的限制，容易出現壓力過多分擔到某些門控專家的現象，造成這些專家所在的機器節點參數更新頻繁，而未發揮作用的專家所在的機器空轉。提出了三個維度的均衡優化，把不同機器上專家的協作屬性融入到loss計算中：

專家維度的均衡，避免有些專家過度勞累，把知識學雜了：

機器維度的均衡，希望處理每個token的6個專家，盡可能分散到不同的機器上：

通信維度的均衡：雖然前面已經做了機器維度的均衡保證，但我們舉一個例子（ep_size=8）：

[tok_0, tok_1, tok_2, ..., tok_n]
算 tok_0 專家所在的機器： 0,1,2,3,5,6
算 tok_1 專家所在的機器： 0,4,2,1,3,7
算 tok_2 專家所在的機器： 0,1,2,3,5,6

這樣仍然不行，雖然滿足了每個token的專家都很分散，但是機器0,1,2,3的使用過于頻繁，4,5,6,7的使用過少。簡單來說，目標2，3聯合起來，理想狀態下是模型參數更新時，專家所在的機器在上方矩陣的行維度最好出現0次或1次，而綜合起來看整個矩陣每個機器出現的次數是整體機器使用量的，這樣才能實現資源利用均衡。