在 iPhone、iPad 和 Mac 上使用 Core ML 加速 Stable Diffusion

TL;DR

Apple 與 Hugging Face 推出了新的 Core ML 優化，具體是 6 位調色板量化和更新的注意力層實現，以讓 Stable Diffusion 在 iPhone、iPad 和 Mac 上運行得更快且記憶體開銷更低。這些改進使得模型權重能夠在精度損失極小的情況下顯著壓縮，從而實現更高效的設備端生成式 AI。

Core ML 優化與 6 位調色板量化

Core ML 利用蘋果設備的 CPU、GPU 和 Neural Engine 在設備上運行機器學習模型。為了提升像 Stable Diffusion 這樣大型模型的效率，Apple 透過 coremltools.optimize 子模組引入了一種名為 6 位調色板量化 的新壓縮技術。

調色板量化的運作方式

調色板量化是一種量化形式，它將模型權重從 16 位浮點表示壓縮為每個參數 6 位。其運作方式類似於電腦圖形中的調色板：定義一個固定數量的顏色（調色板），並將圖像顏色替換為該調色板中最近可用顏色的索引。

記憶體與執行效率

在先前的 Core ML 版本中，壓縮權重在從磁碟載入時會被解壓縮，導致記憶體使用量與未壓縮模型大小相同。更新後的框架現在會在推理逐層進行時即時轉換調色板化權重。此方法更快，因為它減少了記憶體傳輸的量——這是執行的主要瓶頸——而不是增加解壓縮的計算負載。

更新的 Stable Diffusion 實作

Apple 的 ml-stable-diffusion 儲存庫，基於 diffusers 庫，已更新以納入這些新優化：

  • 量化支援： 使用者現在可以使用 --quantize-nbits 標誌將模型量化為 8、6、4 或 2 位。6 位量化被推薦為精度與性能之間的最佳平衡。
  • 注意力層優化： 一種名為 SPLIT_EINSUM_V2 的新方法將查詢序列分割為 512 的塊，以避免創建大型中間張量。此優化可在 Neural Engine 上提升效能 10% 到 30%。

可用的優化模型

模型 未壓縮 6 位調色板量化
Stable Diffusion 1.4 Core ML float16 Core ML 6 位調色板量化
Stable Diffusion 1.5 Core ML float16 Core ML 6 位調色板量化
Stable Diffusion 2 base Core ML float16 Core ML 6 位調色板量化
Stable Diffusion 2.1 base Core ML float16 Core ML 6 位調色板量化

部署與自訂模型轉換

系統需求

要使用 6 位模型，設備必須運行 iOS/iPadOS 17 或 macOS 14 (Sonoma) 的開發版本，因為這些版本包含必要的更新後 Core ML 框架。

轉換自訂模型

開發者可以使用 ml-stable-diffusion 儲存庫來轉換微調或 Dreambooth 模型。該過程包括：

  1. 使用 coremltools 7.0 beta 和 Xcode 15.0 beta。
  2. 使用 --quantize-nbits 6 標誌運行 torch2coreml 轉換腳本。
  3. 在 ORIGINAL 注意力實作（通常在 macOS 上表現更好）和 SPLIT_EINSUM_V2（通常在 iOS 上更快）之間進行選擇。

進階壓縮：訓練時量化

雖然 6 位調色板量化是訓練後壓縮的一個例子，但 coremltools 也引入了訓練時壓縮。這使得開發者可以在同時進行權重壓縮的情況下微調模型。透過使用可微分的權重聚類算法，量化表可以在訓練期間進行優化以最小化損失，從而可能實現 4 位或 2 位壓縮，且其品質降低程度低於訓練後方法。

Sources