apple-aiml-research/ml-ane-transformers

Apple Neural Engine (ANE) Transformers

這是什麼 – ane_transformers 是一個小型開源程式庫，提供針對 Apple 的 Neural Engine (ANE) 優化之 Transformer 模型的參考 PyTorch 實作。它提供兩個套件：

• ane_transformers.reference – 一個獨立、易讀的實作，展示在 ANE 上高效執行所需的底層變更。
• ane_transformers.huggingface – 部分 Hugging Face 模型類別（例如 DistilBert）的即插即用替代品，保留相同 Python API 的同時整合這些最佳化。

為何重要 – Apple 的 A14、A15、M1、M2 等晶片均內建專用神經處理單元。透過重構計算圖（融合運算、將嵌入查表移至 CPU、針對 ANE 的指令集客製化模型），該程式庫聲稱在透過 Core ML 在裝置上執行時，與原生 PyTorch/torch-script 模型相比，可達成 最高 10 倍的延遲降低 與 最高 14 倍的峰值記憶體降低。

如何使用 – README 提供了 Hugging Face distilbert-base-uncased-finetuned-sst-2-english 模型的完整教學：

1. 使用 transformers.AutoModelForSequenceClassification 加載標準模型。
2. 替換為 ane_transformers.huggingface.distilbert.DistilBertForSequenceClassification，並複製原始 state dict。
3. 對輸入進行分詞，使用 torch.jit.trace 追蹤模型，並使用 coremltools.convert 將追蹤後的模型轉換為 Core ML mlprogram 套件。
4. 將生成的 .mlpackage 拖入 Xcode 專案，並使用 Xcode 的效能標籤頁在 iPhone 或 Mac 上驗證加速效果。

安裝 – 該套件已發布於 PyPI：

pip install ane_transformers          # 預先建置的 wheel（最快）
# 或用於開發：
pip install -e .

若在 tokenizers 依賴項上建置失敗，README 指出一個涉及 Rust 編譯器的已知修復方法。

支援的硬體 – 最佳化針對包含 ANE 的 Apple Silicon（iPhone 12/A14 或更新機型、A14+ 的 iPad、M1 或更新的 Mac）進行調校。在舊硬體上執行時，單元測試會發出警告，但生成的 Core ML 模型仍可執行（只是無法獲得宣傳的加速效果）。

限制 / 注意事項

• 目前僅包裝少數 Hugging Face 模型（範例中為 DistilBERT）。擴展至其他架構需額外工作。
• 最佳化後的模型增加一次性的編譯成本；首次載入較慢。
• 606 個運算中有 4 個仍留在 CPU 上（嵌入查表），這是針對此模型大小的有意設計。
• 性能提升取決於序列長度與批次大小；README 指出，對於較長序列（例如 512 個標記，批次 8）可達最高 10 倍的延遲減少。

誰應該關注此專案 – 希望在本地部署 NLP 模型並需要最佳設備內延遲與記憶體佔用的 iOS/macOS 開發者，以及對如何將 Transformer 圖適配至 Apple 的 ANE 感興趣的研究人員。