Pytorch 轉 ONNX 小筆記

基本上這個轉換有兩大類方法, 一類是用官方 tool 去轉, 另一類就是寫個 Python 小程式去做. 原先我都是嘗試後面這路, 但要顧慮的東西很多, 一下修語法, 一下 memory 爆掉, 而是默默出錯時也會轉出 model, 要測試過才知道它的智力有沒有受損?搞得滿累的.

當我再次卡在下面這個檔案限制時, 我就決定換方法了 (悔不當初).

RuntimeError: The serialized model is larger than the 2GiB limit imposed by the protobuf library. Therefore the output file must be a file path, so that the ONNX external data can be written to the same directory. Please specify the output file name.

官方做法其實很簡單, 唯一要顧慮的是 onnx, onnxruntime, onnxruntime_genai 這三個軟體有沒有跟系統衝突? 有沒有跟 NPU tool 衝突 ? 這些搞定就可以了. 用 CPU 也不會轉太久. 這次的障礙是 DeepSeek 跟我講錯指令, 下面這行跑起來找不到 builder.

python -m onnxruntime_genai.builder --model microsoft/phi-2 --precision fp16

我去 onnxruntime 安裝的目錄下找, 確實也沒有對應的程式, 所以我把 Monica 預設的 DeepSeek R1 切到提供第二個意見的 Claude Sonnet V3.7, 它就指出 DeepSeek 的錯誤了, 哈!正確指令如下:

python -m onnxruntime_genai.models.builder \
  --model microsoft/phi-2 \
  --precision fp32 \
  --output ./phi-2-onnx \
  --execution_provider cpu \
  --cache_dir ~/.cache/huggingface \
  --extra_options trust_remote_code=True

轉完之後, 當然要測試一下有沒有問題? 如果發現它答非所問, 應該就是轉錯了. 然而, 我發現 DeepSeek R1 寫的測試程式還是遜了一點, 所以我又讓 Claude 重寫一次.

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
from typing import List, Dict, Optional, Tuple
import time

class Phi2ONNXGenerator:
    def __init__(self, model_path: str, tokenizer_path: str = "microsoft/phi-2"):
        """初始化 Phi-2 ONNX 生成器"""
        # 載入分詞器
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        
        # 設定 ONNX 執行選項以優化效能
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        sess_options.intra_op_num_threads = 4  # 調整為您的 CPU 核心數
        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        
        # 建立推理會話
        self.session = ort.InferenceSession(
            model_path, 
            sess_options=sess_options,
            providers=['CPUExecutionProvider']
        )
        
        # 獲取模型輸入輸出資訊
        self.input_names = [input.name for input in self.session.get_inputs()]
        self.output_names = [output.name for output in self.session.get_outputs()]
        
        # 模型常數
        self.num_layers = 32  # Phi-2 有 32 層注意力層
        self.head_dim = 80    # 每個注意力頭的維度
        self.num_heads = 32   # 注意力頭數量
        
        # 快取字首
        self.key_prefix = 'past_key_values.'
        self.key_suffix = '.key'
        self.value_suffix = '.value'

    def _initialize_kv_cache(self, batch_size: int = 1) -> Dict[str, np.ndarray]:
        """初始化 KV 快取為零張量,使用預分配記憶體"""
        kv_cache = {}
        for i in range(self.num_layers):
            k_name = f'{self.key_prefix}{i}{self.key_suffix}'
            v_name = f'{self.key_prefix}{i}{self.value_suffix}'
            
            # 預分配零張量
            kv_cache[k_name] = np.zeros(
                (batch_size, self.num_heads, 0, self.head_dim), dtype=np.float32
            )
            kv_cache[v_name] = np.zeros(
                (batch_size, self.num_heads, 0, self.head_dim), dtype=np.float32
            )
        return kv_cache

    def _prepare_inputs(self, 
                        input_ids: np.ndarray, 
                        attention_mask: np.ndarray, 
                        kv_cache: Optional[Dict[str, np.ndarray]] = None) -> Dict[str, np.ndarray]:
        """準備模型輸入"""
        inputs = {
            'input_ids': input_ids,
            'attention_mask': attention_mask
        }
        
        # 加入 KV 快取(如果提供)
        if kv_cache:
            inputs.update(kv_cache)
            
        return inputs

    def _update_kv_cache(self, outputs, start_idx: int = 1) -> Dict[str, np.ndarray]:
        """從模型輸出更新 KV 快取"""
        kv_cache = {}
        for i in range(self.num_layers):
            k_idx = start_idx + 2*i
            v_idx = start_idx + 2*i + 1
            
            k_name = f'{self.key_prefix}{i}{self.key_suffix}'
            v_name = f'{self.key_prefix}{i}{self.value_suffix}'
            
            kv_cache[k_name] = outputs[k_idx]
            kv_cache[v_name] = outputs[v_idx]
            
        return kv_cache

    def generate(self, 
                prompt: str, 
                max_new_tokens: int = 100,
                temperature: float = 1.0,
                top_k: int = 50,
                top_p: float = 0.9,
                do_sample: bool = True) -> str:
        """生成文本"""
        start_time = time.time()
        
        # 編碼輸入文本
        encoded_input = self.tokenizer(prompt, return_tensors="np")
        input_ids = encoded_input['input_ids'].astype(np.int64)
        attention_mask = encoded_input['attention_mask'].astype(np.int64)
        
        # 初始化 KV 快取
        kv_cache = self._initialize_kv_cache()
        
        # 初始化輸入
        onnx_inputs = self._prepare_inputs(input_ids, attention_mask, kv_cache)
        
        # 保存原始提示的 token IDs
        prompt_ids = input_ids[0].tolist()
        generated_ids = []
        
        # 逐步生成文本
        for i in range(max_new_tokens):
            # 執行推理
            outputs = self.session.run(None, onnx_inputs)
            
            # 獲取 logits
            logits = outputs[0][:, -1, :]  # [batch, vocab_size]
            
            # 應用溫度
            if temperature > 0:
                logits = logits / temperature
            
            # 選擇下一個 token
            if do_sample:
                # Top-K 過濾
                if top_k > 0:
                    indices_to_remove = logits < np.partition(logits, -top_k, axis=-1)[..., -top_k:][..., :1]
                    logits[indices_to_remove] = -float('Inf')
                
                # Top-p (nucleus) 採樣
                if top_p < 1.0:
                    sorted_logits = np.sort(logits, axis=-1)[:, ::-1]
                    cumulative_probs = np.cumsum(np.exp(sorted_logits) / np.sum(np.exp(sorted_logits), axis=-1, keepdims=True), axis=-1)
                    
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].copy()
                    sorted_indices_to_remove[:, 0] = False
                    
                    # 將索引轉換回原始順序
                    indices_to_remove = np.zeros_like(logits, dtype=bool)
                    for batch_idx in range(logits.shape[0]):
                        indices_to_remove[batch_idx, np.argsort(-logits[batch_idx])[sorted_indices_to_remove[batch_idx]]] = True
                    
                    logits[indices_to_remove] = -float('Inf')
                
                # 計算概率並採樣
                probs = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)
                next_token = np.random.choice(probs.shape[-1], p=probs[0])
            else:
                # 貪婪解碼
                next_token = np.argmax(logits, axis=-1)[0]
            
            # 終止條件
            if next_token == self.tokenizer.eos_token_id:
                break
                
            # 更新生成的 token 列表
            generated_ids.append(int(next_token))
            
            # 更新輸入
            onnx_inputs['input_ids'] = np.array([[next_token]], dtype=np.int64)
            
            # 更新注意力遮罩
            new_attention_mask = np.ones((1, attention_mask.shape[1] + 1), dtype=np.int64)
            new_attention_mask[0, :attention_mask.shape[1]] = attention_mask[0]
            attention_mask = new_attention_mask
            onnx_inputs['attention_mask'] = attention_mask
            
            # 更新 KV 快取
            kv_cache = self._update_kv_cache(outputs)
            onnx_inputs.update(kv_cache)
        
        # 計算生成時間
        generation_time = time.time() - start_time
        tokens_per_second = len(generated_ids) / generation_time if generation_time > 0 else 0
        
        # 解碼並返回生成的文本
        result = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
        
        print(f"生成了 {len(generated_ids)} 個 tokens,耗時 {generation_time:.2f} 秒 ({tokens_per_second:.2f} tokens/秒)")
        
        return result

# 使用範例
if __name__ == "__main__":
    # 初始化生成器
    generator = Phi2ONNXGenerator(
        model_path='./phi-2-onnx/model.onnx',
        tokenizer_path="microsoft/phi-2"
    )
    
    # 生成文本
    prompt = "find all prime numbers below 120"
    result = generator.generate(
        prompt=prompt,
        max_new_tokens=200,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    
    print(f"\n提示:\n{prompt}")
    print(f"\n生成結果:\n{result}")

DeepSeek R1 給的 inference 程式會寫出大致正確但有錯誤的程式 – 邏輯正確, 但引用函數未定義. 我以為 PHI-2 的極限就是這樣了. 想不到 Claude inference 程式寫得好, 答案竟然也跟著好很多 (雖然還有錯)! 在同樣的 model 下也會有顯著的差異, 令我太意外了.

Claude V3.7 Inference 產生的答案:

import numpy as np

def find_primes(n):
    primes = np.arange(2, n)
    for i in range(2, int(np.sqrt(n))+1):
        primes = primes[primes%i!= 0]
    return primes

print(find_primes(120))

DeepSeek-R1 Inference 產生的答案:

import numpy as np

# Define the upper limit
upper_limit = 120

# Create an array of numbers from 2 to the upper limit
numbers = np.arange(2, upper_limit)

# Use the isprime function to find all prime numbers
primes = numbers[np.vectorize(isprime)(numbers)]

print(primes)

DeepSeek 重點分解 – PTX和蒸餾

先前做了一些 DeepSeek 算法上的研讀, 不過其實它的亮點還有很多. 這邊補充一小一大兩個東西. 第一個是 PTX, 第二個是蒸餾.

先前在 “輝達之道" 那篇稍微提到 PTX (Parallel Thread Execution). 在還沒有 CUDA 之前, 輝達就可以使用 Cg, OpenGL 或是 PTX 寫程式. 根據幾篇報導 [1-2] 指出, 這次 DeepSeek 不使用 CUDA, 直接使用 PTX 所以榨出更多的效能.

效能問題只是一個角度, 就好像說我的 code 都是用組合語言寫的, 所以效能更好. 人家可能說你神經病. 但跳過 CUDA 確實不一樣. 很多人認為, 就算大陸做出一個新模型, 效能更好, 還是逃不開輝達的 CUDA, 所以輝達的護城河仍在!像是我相當佩服得美投君 [3] 的新片也是這樣想.

[3] 影片 8’37″

不過我更願意相信, DeepSeek 有意擺脫 CUDA, 而不只是單純為了提升效能. 首先 PTX 類似 Java, 是 just in time 的編譯器 (virtual ISA), 針對不同的硬體可以做二次移植. 其次是 AMD GPU 和華為的 NPU 都支援 DeepSeek [4]. 華為的 Huawei Ascend NPU [5] 是什麼概念呢? 它可以中國產商不買輝達 GPU, 美國掐不住它的脖子 [6-7].

其次談一下蒸餾 (distillation). 我們知道這是一個老師教學生的演算法, 大模型教小模型, 小模型甚至能青出於藍, 但計算量就省下來了! DeepSeek 是一個大模型, 就算 MOE 等方法可以讓它只激活部分參數, 那還是很巨大啊!

但讓世人震撼的另外一面是 DeepSeek R1! 現在如果問 DeepSeek 關於它自己的技術, 它會有點故意誤導, 不知道是不是政治干擾? 根據 HuggingFace 上的說明: “DeepSeek-R1-Zero & DeepSeek-R1 are trained based on DeepSeek-V3-Base." [8] 但 DeepSeek 自己倒是說 “R1 可能是 DeepSeek 的早期版本或基础版本", 連自己的身世都胡扯, 哈!

回到正題, R1 比 V3 更強大 [8], 它對標的是 OpenAI o1 – 1217. V3 是最右邊偏矮的那一個. AI 要做得好, 除了硬體和演算法, 就是要靠好的教材. V3 反映了演算法, R1 彰顯的就是好的教材.

R1 有兩個版本, DeepSeek R1 和 DeepSeek R1- Zero. 在官網講得有點不清楚, 我先引用曲博的影片 [9], 再補充名詞解釋.

R1-Zero = V3 + RL => 無盡的重複、可讀性差、語言混雜 => 有缺點.

R1 參考下圖.

[9] 17’19″

其中:

  • RL = large-scale reinforcement learning, SFT = supervisor fine tuning,
  • cold start = 通常是指參數是未訓練過、隨機的. 但 SFT 後再 cold start 有點怪怪的. 這部分還不理解.
  • GRPO (group relative Policy Optimization) = 群內相對評比. 也就是標準答案不從外面給, 而是自己比較哪個答案好? 例如寫兩段 code, 誰的效率高自己知道.

接下來就是蒸餾的部分, DeepSeek 推出了它當老師教小模型的版本. 其中交 Qwen2.5-32B 就已經很厲害, 教 Llama 70B 的部分, 在下表 [8] 的比試幾乎全勝! 只有 CodeForces rating 這項還輸 o1-mini 而已.

這表示什麼呢? 這說明就算你的系統還跑不了 V3 或是更優化的 R1, 只要用它去教小模型, 小模型也堪用. 像是它教出來的 Qwen 只有 1.5B, 好幾項測試都還贏 OpenAI 的 Claude 3.5 Sonnet, 這個還要花錢買耶 (我! QQ). Qwen 卻是 open source.

所以我感覺 DeepSeek 的發布不只是火力展示, 它的目的是要從算法、數據 (能教人是好老師) 兩方面展示不輸老美的優勢. 強強相爭, 以後 AI 會進步地更快吧! 只要不消滅人類都是好事!

[後記]

寫完之後, 發現老高也來評論了 [11]. 他講到一個值得補充的地方. ChatGPT 4o 對標 DeepSeek V3, 訴求全能. ChatGPT 01 對標 DeepSeek R1, 訴求推理過程 (chain of thought).

[REF]

  1. DeepSeek 绕开 CUDA 垄断,针对英伟达 PTX 进行优化实现最大性能,英伟达护城河还在吗
  2. https://technews.tw/2025/01/29/deepseek-bypass-cuda-and-using-ptx-for-better-optimization/
  3. https://www.youtube.com/watch?v=81cbmeXTQcg
  4. https://huggingface.co/deepseek-ai/DeepSeek-V3
  5. https://medium.com/huawei-developers/world-of-huawei-ascend-future-with-npus-5843c18993f3
  6. https://blog.csdn.net/qq_54958500/article/details/144064251
  7. https://udn.com/news/story/7333/8022387
  8. https://huggingface.co/deepseek-ai/DeepSeek-R1
  9. https://www.youtube.com/watch?v=spoPf8CjjBo
  10. https://www.bnext.com.tw/article/79507/claude-3.5-sonnet-best-ai?
  11. https://www.youtube.com/watch?v=uKBI1Ea8VO0

DeepSeek 重點分解 – MTP 小整理

先前分析了 V2 的主力武器, 但 V3 還是比 V2 厲害一截. 所以來談一下 V3 新增的 multi-token prediction (MTP). 雖然 V3 還有厲害的 pre-train 和 fine tune, 但那部份無法用數學或是圖形表示, 只好略過.

本圖取材自 https://github.com/deepseek-ai/DeepSeek-V3

顧名思義, multi-token prediction 就是一次預測好幾個 token. 問題來了, 究竟是一次預測好幾個 tokens (下圖左)?還是預測完一個繼續預測下一個 (下圖中)?還是一次預測好幾個又連續預測好幾步 (下圖右)?

本圖取材自 https://arxiv.org/html/2410.03132v3

其實眼尖一點就可以看到上圖右 (Ours) 一定是該論文認為最好的. 但是這種暴力美學好像跟 DeepSeek 省吃儉用的調性不合. 我們來看看下圖的 DeepSeek V3 架構又是怎做的?

本圖取材自 https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf

我們輸入的 token 是上圖下方的 input t1, t2, t3,… 這些 token, 預測的輸出是上圖上方的 t2, t3, t4….等. 如果我們放置愈多的 MTP module, 就可以預測更深的深度 (Depth). 例如 D = 2, 就是用 t1 預測 t2 和 t3, 用 t2 預測 t3 和 t4, 依此類推. D 愈大, 預測的深度就愈長, 因此它符合 MTP in a single step.

可是一次要預測好幾個 token, 計算不會很大嗎?PDF 提到, 圖裡面的 embedding layer, Output head 雖然畫了好幾個, 但他們都是公用的 (shared).

當我們預測 t2時, 看向上面圖中的 MTP module 1 . 它有且只有兩個輸入, 一個是 t2 經過 embedding, 一個是 t1 在 main Model 的產出物.

至於預測 t3時, 看向上面圖中的 MTP module 2. 它有且只有兩個輸入, 一個是 t3 經過 embedding, 一個是 t2 在 MTP Module 1 的產出物.

因此 D 愈大, 計算的確愈多, 並且有額外的 latency. 但是額外增加的幅度並不等效於再重複一次 main model. 而且這樣有一個好處, 就是 training 時可以好好吸收長距離依賴 (long-range dependencies), 因為它每個預測都以過去歷史為本, 不會即興創作.

至於 multi-token prediction in multiple steps (Autoregressive Chunking) 的方法, DeepSeek 認為在 inference 時確實比較好. 但是在 training 的時候, MTP in one step (Action Chunking) 比較能控制因果關係和長距離依賴.

好!不愛數學的可以停在這裡. 接下來是講公式的部份. 其實跟上面一模一樣, 只是 step 3 增加了以logits 做 softmax() 去字典找字.

Step 1: Combining Representations

For the 𝑖-th input token 𝑡𝑖 at the 𝑘th prediction depth:

  1. The representation of the 𝑖th token at the (𝑘−1)th depth, denoted as h𝑘−1i∈ ℝ^𝑑, is taken. If 𝑘 = 1, h𝑘−1𝑖 is the representation provided by the main model.
  2. The embedding of the (𝑖+𝑘)th token, Emb(𝑡𝑖+𝑘) ∈ ℝ^𝑑, is computed using the shared embedding layer.
  3. Both h𝑘−1𝑖 and Emb(𝑡𝑖+𝑘) are normalized using RMSNorm (Root Mean Square Normalization).
  4. The normalized representations are concatenated ([·; ·]) and linearly projected using the projection matrix𝑀𝑘 :
    • h′ki= Mk[RMSNorm(hk−1i); RMSNorm(Emb(ti+k))].
    • Here, h′𝑘𝑖 is the combined representation that serves as the input to the Transformer block at the 𝑘th depth.

Step 2: Transformer Block

The combined representation h′𝑘𝑖 is passed through the 𝑘th Transformer block (TRM𝑘(·)): h𝑘1:𝑇−𝑘 = TRM𝑘(h′𝑘1:𝑇−𝑘).

This produces the output representation h𝑘𝑖 for the 𝑖th token at the 𝑘th depth. The slicing operation 1:𝑇−𝑘 ensures that the sequence length is adjusted appropriately for each prediction depth.


Step 3: Output Head

The output representation h𝑘𝑖 is passed through the shared output head (OutHead(·)), which:

  1. Linearly maps h𝑘𝑖 to logits.
  2. Applies the Softmax function to compute the probability distribution over the vocabulary:𝑃𝑘𝑖+𝑘+1 = OutHead(h𝑘𝑖).
    • Here, 𝑃𝑘𝑖+𝑘+1 ∈ ℝ^𝑉 represents the probability distribution for the (𝑖+𝑘+1)th token, where 𝑉 is the vocabulary size.

最後一個重點來了. DeepSeek 只有在 training 的時候使用 one step MTP. 在 inference 的時候, 用的演算法又有不同. “We can also repurpose these MTP modules for speculative decoding (預言家, 投機演算法) [2] to further improve the generation latency."[1]

Training 的 loss function 計算也給出來了. 首先, 針對每個 depth (k) 都做計算, P 就是上面的 P. 最後把不同深度的 loss function 取平均值.

其中 𝜆 當然就是 weighting factor, 或是以前電力機械教授老包所說的 “那麼大”. γ𝜆 = “柑仔那麼大", 是我對這堂課最深的印象.

[REF]

  1. https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf
  2. https://hackmd.io/@shaoeChen/rJESTVr40

DeepSeek 重點分解 – MOE 小整理

DeepSeek 另外一個重點是 MOE (mixure of Expert). 從抽象觀念來理解, 就是說: 既然 Model 裡面有很多個 Expert, 問問題的時候, 只要把特定的專家叫起來回答就可以省能量, 加快反應速度.

LLM 的專家能力體現在 feed forward network (FFN) 這個部分. 以我的理解來說, 過去我們認為的類神經網路的確可以存儲知識, 只是差在沒有 transformer 來理解和表達語意. 人類的大腦也是分化為很多特定區域 [1], 因此動腦的時候的確不需要火力全開. 把 model 分為多個 expert 是很直覺的事.

好, 那麼誰決定選哪一些 expert 起來工作呢? 我們把它叫做 routed expert. 這個路由專家根據 input 把負責 sub-network 的串起來. 不過, 萬一串錯了, 雞同鴨講怎麼辦? 沒關係, 還有 shared expert, 它其實是個通才.

在下面的式子中, ut 代表第 t 個 token 對 FFN 的輸入. FFN(s) 代表 shared expert FFN, FFN(r) 代表 routed expet FFN, 看起來雖然有一點小複雜, 但其實 ht‘就是原始輸入 + routed FFN + shared FFN.

要加哪幾個 routed expert, 取決於 gi,t, 這個是什麼呢? 我們先解釋一下 e, 其他就很好理解了. e 表示 expert 的質心 (the centroid of the routed expert), 可以理解為 expert 的 “代表矩陣". 可能就像我們以前做 clustering, 每個 cluster 都要有一個中心 mean, 這樣才能分群.

當輸入 uT和 e 做矩陣相乘, 直覺可以知道它在判斷輸入 u 和 expert e 是否相似? si,t 就是這個輸出的 softmax 正規化. 既然有值就可以比大小, 最大的 k 個 (top k) si,t 個專家, gi,t不等於 0 就表示會被召喚.

到目前為止, 一切都很直覺. DeepSeek 最厲害在哪裡呢? 除了它把專家切得很細, 還要求挑選的專家要跨 device 做 load-balance. 當 device 數 >= 3, 效果會比較好. 為什麼要搞這個呢? 我認為就是美國不賣它厲害的 GPU 嘛! 每個 GPU 的 memory 都有限制, 所以當然產生了 3 個臭皮匠勝過一個諸葛亮的架構!

為了防止只有某幾個臭皮匠過勞, 或是臭皮匠集中在一起產生工作熱區, 或是遠距溝通但產生通訊熱區. DeepSeek 設計了一些 load balance 的機制. 因此我認為老美的管制是有效的. 如果有水冷的高級 GPU array 和 NPU, 基本上錢砸下去就有了. 因此相關的段落可以跳過. 但是 token dropping [2] 的部分值得一提.

即使上面已經盡力去做 load balance, 還是有可能某些 device 在 “訓練時" 超過運算 budget. 此時就把最沒有親和力 (affinity score, 與上下文相關性) 的 token 丟掉, 丟到平衡為止. 同時又保證至少 training sequence 中的 10% 絕對不能丟.

we drop tokens with the lowest affinity scores on each device until reaching the computational budget. In addition, we ensure that the tokens belonging to approximately 10% of the training sequences will never be dropped. 

有了 MLA 和 MOE, DeepSeek 的兩大武器都稍微提到了. 隱藏在演算法後面的, 算是人海戰術清理訓練資料吧? Microsoft 的 PHI-2 系列就證明過: 只學有用的東西就好, 叫 AI 學一堆垃圾, 卻硬要訓練到收斂根本不環保. 下面稍微列出該論文 [3] 對 pre-trained 下的功夫, 我對於 DeepSeek 的追蹤就暫時告一段落.

While maintaining the same data processing stages as for DeepSeek 67B 
(DeepSeek-AI, 2024), we extend the amount of data and elevate the data quality. In order to enlarge our pre-training corpus, we explore the potential of the internet data and optimize our cleaning processes, thus recovering a large amount of mistakenly deleted data. Moreover, we incorporate more Chinese data, aiming to better leverage the corpus available on the Chinese internet. In addition to the amount of data, we also focus on the data quality.

We enrich our pre-training corpus with high-quality data from various sources, and meanwhile improve the quality-based filtering algorithm. The improved algorithm ensures that a large amount of non-beneficial data will be removed, while the valuable data will be mostly retained. In addition, we filter out the contentious content from our pre-training corpus to mitigate the data bias introduced from specific regional cultures. A detailed discussion about the influence of this filtering strategy is presented in Appendix E.

[REF]

  1. 我讀 «大腦超載時代的思考學» – 2
  2. C. Riquelme, J. Puigcerver, B. Mustafa, M. Neumann, R. Jenatton, A. S. Pinto, D. Keysers, and N. Houlsby. Scaling vision with sparse mixture of experts.In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, pages 8583–8595
  3. https://arxiv.org/html/2405.04434v5

DeepSeek 重點分解 – MLA 小整理

DeepSeek [1] 重點之一是 MLA (Multi-head Latent Attention) . 它可以單獨使用. 解釋它時, 若和 MHA (Multi-Head Attention) 對比會更好理解, 所以先回顧一下 MHA.

1. 原始的注意力計算

一般的注意力機制中, 主要有 q、k、v 三個矩陣. 分別代表 query, key, 和 value. W 是權重矩陣, h 是 hidden state. 下標 t 表示第 t 個 token.

加入多頭機制後: d 是 embedded dimension, nh = attension head 數, , dh 是每個 head 的 dimension. 其中 d = nh * dh , j 是從第一個到第 t 個 token 的 index, i 是第一到第 nh 個 head 的 index.

Attention o 當然也分為第幾個頭的第幾個 token. 故表示為 ot,i. 同樣多頭都用 o 權重矩陣 Wo 轉出 output ut.

一般認為 k, v 的值太多, 是造成計算量和記憶體過多的元凶. 但是不記住這些東西, transformer 就發揮不出造句的能力. 科學家想了各種解法想要簡化 k,v. 但是操作不好就會降低 LLM 的性能. DeepSeek 使用的 MLA 看起來可以兩全其美.

2. MLA 中的計算流程

2.1 Low-Rank Key-Value Joint Compression

在 MLA 中, 首先對 hidden state 壓縮. 從前面的 MHA 的段落可以得知, q, k, v 共用 hidden state 但不共用權重. 那麼我們壓縮 h 再還原就可以節省參數量了. c 矩陣由 ht 乘 WDKV 而來. 顧名思義, D 代表 down projection, kv 矩陣意義跟先前相同. 做完下投影再做上投影, 理解為壓縮解壓縮即可. 所以 WUK 還原 k, WUV 還原 v. 其中 U 就代表 up projection.

我們已經知道這是濃縮再還原的果汁. 為了確定風味不變太多. 損失的部分要要別的方法補回來. 甚至 DeepSeek 為了減少 active node, 連你問的問題 q = query 都壓縮了. 這個猛! 表示我跟它說 “請、謝謝、對不起" 都是多餘的.

2.2 Decoupled Rotary Position Embedding

上面那招還不是全貌. 但是要講第二招就要先講 RoPE (Rotary Position Embedding) [2]. 我們知道原本 transformer 就要記錄 token 的相對位置關係. 畢竟"你愛我" 跟 “我愛你" 是兩碼子事. RoPE 這個演算法就是用來把位置編碼專屬的維度給省了, 但它結 “繩" 記事, 還記得相對位置.

不過這招和 2.1 壓縮那招有衝突, 壓縮完再解壓就不符合交換律. 所以又衍生出 decouple 的輔助算法, 再浪費一點空間. 為 RoPE 額外產生的 dimension 為 dhR. 多出來的 qtR 和 ktR 加在原來的矩陣後面. 式子大致上都一樣.

3. 效能比較

假如我們比對 MHA 和 MLA, 就會發現它的 KV cache 比較少, 而且實測效果更好. 至於 GQA 和 MQA 是來陪榜的. 用論文 [1] 中的圖解帶過.

[REF]

  1. https://arxiv.org/abs/2405.04434
  2. . J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu.Roformer: Enhanced transformer with rotary position embedding.Neurocomputing, 568:127063, 2024.