在 Core ML 上執行 Mistral 7B

WWDC 24 的新 Core ML 功能

Swift Tensor（MLTensor）

全新的 MLTensor 類型在 Swift 中提供了多維資料結構的高階抽象，模仿 numpy 陣列或 torch 張量的功能。它取代了先前需要以不透明指標存取底層儲存空間的 MLMultiArray 或 MLShapedArray 的低階操作。MLTensor 內建了如 softmax 等運算，簡化了語言模型的前處理與後處理流程。

有狀態緩衝區

傳統上，Core ML 模型以無狀態函式運作。有狀態緩衝區允許模型在 GPU 上預留一塊記憶體，以在多次迭代間維持狀態。對於大型語言模型（LLM）而言，這用於實作鍵值（KV）快取。將 KV 快取保留在 GPU 上，可避免在每次產生 token 時在 CPU 與 GPU 之間傳送大量有狀態資料的開銷，透過減少記憶體頻寬瓶頸顯著提升效能。

區塊式量化

為了在不顯著降低品質的前提下縮減模型大小，Core ML 現在支援區塊式量化。此技術為同一張量的不同區域建立多個查找表（LUT），而非為整個張量使用單一表格。這使模型能壓縮至 4 位元精度，較 float32 減少 8 倍大小，較 float16 減少 4 倍。

多功能支援

多功能支援允許開發者將 LoRA（低秩適應）適配器封裝進生成模型。透過交換少量額外參數（適配器），即可讓單一基礎模型用於不同任務或風格，無需重新載入整個模型。

將 Mistral 7B 轉換為 Core ML

要有效執行 Mistral 7B，需要自訂的注意力實作，預先分配完整的 KV 快取緩衝區並在原位更新，以符合 Core ML 的有狀態緩衝區需求。

追蹤與轉換流程

  1. Tracing: 模型使用已修補的 StatefulMistralForCausalLM 實作載入，並透過 torch.jit.trace 以範例輸入進行追蹤。
  2. Input Definition: 透過 coremltools 使用範圍維度，使輸入序列能從單一 token 擴展至最大 2048 的上下文長度。
  3. State Preparation: keyCache 與 valueCache 緩衝區使用 ct.StateType 定義，以確保它們被轉換為有狀態的 Core ML 緩衝區。
  4. Conversion: 模型以 iOS 18 或 macOS 15 為最低部署目標轉換為 Core ML，以存取新的有狀態功能。

模型壓縮

使用 OpLinearQuantizerConfig，模型以 4 位元線性對稱量化（區塊大小 32）進行壓縮。最終的 mlpackage 大小從約 14GB（float16）縮減至約 3.8GB。

執行與實作

使用者可以透過 swift-transformers 倉庫的 preview 分支執行已轉換的 Mistral 7B 模型。

使用 Swift 執行

preview 分支整合了完整的 MLTensor 支援以及 Swift 版的有狀態 API。可透過指令列執行推論：

git clone -b preview https://github.com/huggingface/swift-transformers
swift run transformers "Prompt text" --max-length 128 Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage

使用 Python 執行

亦可使用 Python 透過 coremltools 進行推論：

python3 generate.py Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage --prompt "Prompt text"

未來路線圖

Hugging Face 計畫將這些實驗性方法整合至 exporters，這是一個用於將 transformers 模型轉換為 Core ML 的 Python 工具。此外，他們也希望透過探索 OpenELM 或 DCLM 等較小模型，優化 iPhone 上 Apple Neural Engine（ANE）的效能，因為目前的 Mistral 7B 實作主要針對 Mac GPU 進行最佳化。

Sources