|
DeepSeek最近憑借一己之力擊落Nvidia神話，從此硬件不再是AI的決定性條件，低成本訓練LLM模型成為現實。之所以能夠實現低成本訓練大模型，主要歸功于DeepSeek底層采用的MLA和MoE結構，MoE結構我們之前已經介紹過，相比于訓練一個統一的專業大模型，MoE通過訓練N的專家模型，能夠更高效的實現推理任務。 
MLA（Multi-head Latent Attention）結構是DeepSeek的核心創新之一，旨在優化Transformer模型中的注意力機制，以減少推理過程中的KV緩存（鍵值緩存）并提高計算效率。其核心原理如下： 1. MLA的核心思想：低秩聯合壓縮MLA的核心思想是對鍵（Key）和值（Value）進行低秩聯合壓縮。具體來說，傳統的Transformer模型在計算注意力時，需要分別計算和存儲每個頭的鍵（K）和值（V），這會導致大量的內存消耗和計算開銷。而MLA通過引入一個低秩向量（c），將K和V的存儲和計算過程進行壓縮。具體步驟如下： - 計算低秩向量c：首先計算一個低秩向量c，其維度遠小于K和V的維度。
- 轉換回K和V：然后通過兩個權重矩陣（W）將c轉換回K和V。
這樣，在推理階段，只需要存儲低秩向量c，而不是完整的K和V，從而顯著減少了KV緩存的大小。 2. 權重矩陣的合并與吸收MLA的另一大特點是權重矩陣的合并與吸收。傳統的MHA（Multi-head Attention）模型中，每個頭都有獨立的Q、K、V投影矩陣，而在MLA中，這些矩陣被合并和吸收，以減少參數量和計算量。具體來說： - Q的投影矩陣：Q的投影矩陣可以與后續的投影計算合并。
- K和V的投影矩陣：K和V的投影矩陣也可以進行類似的合并操作。
通過這種合并，MLA在存儲和計算上都優于傳統的MHA。例如，KV緩存的大小可以從 減少到 ，其中h是head的數量。 3. 解耦的旋轉位置編碼（RoPE）MLA還對旋轉位置編碼（RoPE）進行了優化。RoPE對K和Q的位置敏感，但由于低秩向量c改變了位置，因此無法直接使用原有的位置編碼。為了解決這個問題，MLA采用了以下方法： - 使用額外的多頭Q和共享的K：引入額外的多頭Q和共享的K來攜帶旋轉位置編碼。
這種解耦的方法確保了位置編碼的正確性，同時不影響MLA的低秩壓縮和權重合并。 4. MLA的計算流程MLA的計算流程可以概括如下： - 低秩壓縮：Q和K路徑中分別進行低秩壓縮，生成latent / low-rank部分。
- 解耦RoPE：K路徑中的一部分進行解耦的RoPE變換。
- 緩存與合并：推理階段，壓縮后的KV矩陣和位置編碼后的矩陣被緩存，并進行合并吸收。
5. MLA的優勢- 減少KV緩存：通過低秩聯合壓縮，MLA顯著減少了KV緩存的大小。
- 降低計算成本：權重矩陣的合并與吸收減少了參數量和計算量。
- 提高推理效率：在相同的硬件條件下，MLA可以實現更長的上下文長度或更大的batch size，從而提高推理速度或吞吐量。
6. MLA與MQA的比較MLA的計算過程在某些方面與MQA（Multi-query Attention）類似，但MLA在QK維度上進行了RoPE變換，并且V與未施加RoPE的K共享激活值，這使得MLA在保持高效的同時，仍然能夠保持較好的模型性能。 QuantML-Qlib實現MoE之前已經集成在我們的框架之中，詳見： QuantML-Qlib開發版 | MoE混合專家系統用于提升Transformer表現【附代碼】
本次我們將MLA結構也集成進框架之中，代碼路徑：examples/benchmarks/DeepSeek 
其中MLA結構的代碼為：
class MultiHeadLatentAttention(nn.Module): def __init__(self, config: Config): super().__init__()
assert config.v_head_dim is not None , f"v_head_dim is not defined {config.v_head_dim=}" assert config.q_lora_rank is not None , f"q_lora_rank is not defined {config.q_lora_rank=}" assert config.kv_lora_rank is not None , f"kv_lora_rank is not defined {config.kv_lora_rank=}" assert config.rope_head_dim is not None , f"rope_head_dim is not defined {config.rope_head_dim=}"
self.config = config
self.dim = config.d_model self.num_heads = config.num_heads self.v_head_dim = config.v_head_dim
self.nope_head_dim = config.nope_head_dim self.rope_head_dim = config.rope_head_dim
self.q_lora_rank = config.q_lora_rank self.kv_lora_rank = config.kv_lora_rank
self.dropout = config.dropout
self.value_dim = self.num_heads * self.v_head_dim
# this is dims between wQ and wK self.nope_dim = self.num_heads * self.nope_head_dim self.rope_dim = self.num_heads * self.rope_head_dim
# query compression self.compress_q_linear = nn.Linear(self.dim, self.q_lora_rank, bias=False) # W_DQ self.decompress_q_nope = nn.Linear(self.q_lora_rank, self.nope_dim, bias=False) self.decompress_q_rope = nn.Linear(self.q_lora_rank, self.rope_dim, bias=False) self.q_norm = RMSNorm(dim=self.q_lora_rank) # key and value compression self.compress_kv_linear = nn.Linear(self.dim, self.kv_lora_rank, bias=False) # W_DKV self.decompress_k_nope = nn.Linear(self.kv_lora_rank, self.nope_dim, bias=False) self.decompress_v_linear = nn.Linear(self.kv_lora_rank, self.value_dim, bias=False) self.kv_norm = RMSNorm(dim=self.kv_lora_rank)
self.k_rope_linear = nn.Linear(self.dim, self.rope_head_dim , bias=False) # self.rope_norm = RMSNorm(self.rope_dim) # not in deepseekv2
self.proj = nn.Linear(self.value_dim , self.dim, bias=False) self.res_dropout = nn.Dropout(p=config.dropout)
def forward(self, x: Tensor,mask: torch.Tensor, freqs_cis: Tensor): batch_size, seq_len, _ = x.shape
compressed_q = self.compress_q_linear(x) norm_q = self.q_norm(compressed_q) query_nope:Tensor = self.decompress_q_nope(norm_q) query_rope:Tensor = self.decompress_q_rope(norm_q)
compressed_kv = self.compress_kv_linear(x) norm_kv = self.kv_norm(compressed_kv) key_nope: Tensor = self.decompress_k_nope(norm_kv) value: Tensor = self.decompress_v_linear(norm_kv)
key_rope:Tensor = self.k_rope_linear(x) # norm_rope = self.rope_norm(key_rope)
query_nope = query_nope.view(batch_size, seq_len, self.num_heads, self.nope_head_dim).transpose(1,2) query_rope = query_rope.view(batch_size, seq_len, self.num_heads, self.rope_head_dim).transpose(1,2)
key_rope = key_rope.view(batch_size, seq_len, 1, self.rope_head_dim).transpose(1,2) key_nope = key_nope.view(batch_size, seq_len, self.num_heads, self.nope_head_dim).transpose(1,2)
value = value.view(batch_size, seq_len, self.num_heads, self.v_head_dim).transpose(1,2)
# *** the line that fixes MLA :) *** key_rope = key_rope/self.num_heads
q_rope,k_rope = apply_rope(query_rope,key_rope, cis=freqs_cis)
q_recombined = torch.empty((batch_size,self.num_heads,seq_len, self.rope_head_dim + self.nope_head_dim), device=x.device) k_recombined = torch.empty((batch_size, self.num_heads, seq_len, self.rope_head_dim + self.nope_head_dim), device=x.device)
q_recombined[:,:,:,:self.nope_head_dim] = query_nope q_recombined[:,:,:,self.nope_head_dim:] = q_rope
# k_rope = torch.repeat_interleave(k_rope, self.num_heads, dim=1) # >> you dont need to do this << # ?? broadcasting will do replication krope to all heads automagically k_recombined[:,:,:,:self.nope_head_dim] = key_nope k_recombined[:,:,:,self.nope_head_dim:] = k_rope
output = F.scaled_dot_product_attention(q_recombined, k_recombined, value, is_causal=True, dropout_p=self.dropout)
output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.num_heads * self.v_head_dim)
output = self.proj(output) output = self.res_dropout(output) return output
簡單測試了一下，500內的IC為0.025，歡迎繼續測試優化代碼結果。
|
