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)

ONNX 的 version 小筆記

ONNX (Open Neural Network Exchange) 是一種 AI 檔案交換標準 最近在 porting DeepSeek 到 NPU 的過程中, 卡住了一下, 所以特別記錄下來.

首先 ONNX 有兩個主要的 domain: ai.onnx 和 ai.onnx.ml. 雖然他們都是為 AI 而生, 但ai.onnx domain 的 operator (=運算子, 算子, 運算符) 更加常用, 像是加減、乘加這種很基本的運算, 或是比較複雜的 convolution, activation function, 以及 DNN 有關的 RNN, dropput…等等.

ai.onnx.ml 的算子適用於傳統的機器學習, 包括非深度學習的算法, 像是 one hot encoder, 還有 SVM, imputer…等等. 至於其他的 domain, 屬於客製化擴展, 官方不維護 [1]. 包括: ai.onnx.training, com.microsoft, com.ms.internal….等等.

接下來也很妙. ONNX 進版後, 雖然看起來也是以 1.0, 1.1,….這樣長大, 但是官方又出了個 IR version, 用來表示是否有 breaking change. 所以顯然很多版 version 會對應到同一版 IR version.

“The IR format is versioned using simple numbers, which MUST be monotonically increasing. Breaking changes to the format or semantics of the ONNX specification require an increment of the version. Non-breaking changes to the IR format do not require changing the version number."

光是這樣也就罷了, 由於還有 ai.onnx 和 ai.onnx.ml 兩個 domain 之分. 如果對某個 domain 是 breaking change, 但對另外一個 domain 不是呢? 所以對於不同的 domain, 又有 Opset 的 version.

“ONNX uses operator sets to group together immutable operator specifications. An operator set represents a specific version of a domain, indicated by a pair (domain, version). This represents the set of all operators belonging to the specified domain with the specified version (referred to as the opset_version). When the inventory of a given operator set changes either by addition, removal, or a change in semantics of a contained operator, its version MUST increase."

因此我們可以得到這樣的表格 – Released Versions

ONNX versionIR versionOpset version ai.onnxOpset version ai.onnx.mlOpset version ai.onnx.training
1.0311
1.1351
1.1.2361
1.2371
1.3381
1.4.1491
中間省略, 詳情參考 [1]
1.14.191931
1.15.092041
1.16.0102151
1.17.0102251

當我試圖將 ONNX INT4 的 DeepSeek 翻成 NPU 指令時, converter 跑出下面這個 error, 表示在 domain version 14 當中找不到 SimplifiedLayerNormalization 這個算子. 而我們安裝的 ONNX 通常是 ai.onnx domain. 既然在預設的 domain 裡面找不到, 我們就要檢查一下從 huggingface 抓來的 DeepSeek 是哪個 domain 的 ONNX?

File "/local/workspace/hailo_virtualenv/lib/python3.10/site-packages/onnx/checker.py", line 163, in check_model
    C.check_model_path(
onnx.onnx_cpp2py_export.checker.ValidationError: No Op registered for SimplifiedLayerNormalization with domain_version of 14
 
==> Context: Bad node spec for node. Name: /model/layers.0/input_layernorm/LayerNorm OpType: SimplifiedLayerNormalization

寫個小程式檢查 node 的 domain:

import onnx
model = onnx.load("model.onnx")

# 查看所有算子的 domain
for node in model.graph.node:
    print(f"Op: {node.op_type}, Domain: {node.domain}")

結果悲劇了… 大部分都是 com.microsoft domain, 難怪會出錯. 經過測試, 這個 int4 版本除了回答品質不太穩之外, 還會無限迴圈重複自己的話, 所以暫且就不 debug 下去了, 僅記錄下 ONNX version 的複雜性.

[REF]

  1. https://github.com/onnx/onnx/blob/main/docs/Versioning.md