sophgo/tpu-mlir

Machine learning compiler based on MLIR for Sophgo TPU.

解決的問題

TPU-MLIR 解決了將來自各種主流框架的預訓練神經網路部署到 TPU 硬體上的問題。它彌合了高層模型定義（如 PyTorch 或 ONNX）與 TPU 晶片高效執行所需的特定二進位格式（bmodel）之間的鴻溝。

工作原理

該專案使用基於 MLIR（多層中間表示）的編譯流水線，透過以下幾個階段轉換模型：

  1. 前端匯入：從 PyTorch、ONNX、TFLite、Caffe 或 HuggingFace 匯入模型並轉換為統一的 MLIR 表示。
  2. 降低（Lowering）：編譯器透過模式重寫和記憶體規劃，將高層表示轉換為 TPU 特有的方言（Top 和 Tpu）。
  3. 量化：應用各種量化技術（F32、BF16、F16、INT8）和校準，以優化硬體效率。
  4. 部署：最終輸出是可以在目標 TPU 處理器上運行的 bmodel 檔案。

適用對象

專為需要將深度學習模型（包括大型語言模型和視覺模型）部署到 Sophgo TPU 硬體的 AI 工程師和開發人員設計。

亮點

  • 廣泛的框架支援：相容於 PyTorch、ONNX、TFLite、Caffe 和 HuggingFace。
  • LLM 優化：針對 HuggingFace LLM 的專用單次轉換，支援歷史 KV 快取和分塊推理。
  • 全面的量化：支援對稱/非對稱 INT8、F16、BF16，以及 AWQ、GPTQ 和 AutoRound 的透傳。
  • 生產級工具鏈：包括通用推理執行器 (model_runner)、bmodel 檢查工具 (model_tool) 和精度驗證實用程式。