釋出 Swift Transformers：在蘋果裝置上執行裝置端 LLM

釋出於 2023 年 8 月 8 日

我非常尊敬 iOS/Mac 開發者。我從 2007 年開始為 iPhone 編寫應用程式，當時甚至還沒有 API 或文件。新裝置在約束空間上採納了一些不熟悉的決策，結合了功耗、螢幕空間、UI 習慣、網路訪問、永續性和延遲等因素，這與我們之前習慣的有所不同。然而，這個社群很快就設法創造出與新正規化融為一體的頂級應用程式。

我相信機器學習（ML）是一種構建軟體的新方式，而且我知道許多 Swift 開發者希望將 AI 功能整合到他們的應用中。ML 生態系統已經非常成熟，有數千個模型可以解決各種各樣的問題。此外，大型語言模型（LLM）最近已成為近乎通用的工具——只要我們能將任務建模為處理文字或類文字資料，它們就可以適應新的領域。我們正在見證計算歷史上的一個決定性時刻，LLM 正走出研究實驗室，成為每個人的計算工具。

然而，在應用程式中使用像 Llama 這樣的 LLM 模型涉及多個任務，許多人需要獨自面對和解決。我們一直在探索這個領域，並希望與社群繼續合作。我們的目標是建立一套工具和構建模組，幫助開發者更快地進行構建。

今天，我們釋出這篇指南，介紹在 Mac 上使用 Core ML 執行像 Llama 2 這樣的模型所需的步驟。我們還發布了 alpha 版的庫和工具，以支援開發者的這一旅程。我們呼籲所有對 ML 感興趣的 Swift 開發者——這是否意味著*所有* Swift 開發者？——透過提交 PR、報告 bug 或提出意見來共同改進這一切。

衝呀！

今日釋出

swift-transformers，一個正在開發中的 Swift 包，旨在 Swift 中實現類似 transformers 的 API，專注於文字生成。它是 swift-coreml-transformers 的演進版本，目標更廣泛：Hub 整合、支援任意分詞器以及可插拔模型。
swift-chat，一個演示如何使用該包的簡單應用程式。
exporters 的更新版本，一個用於 transformers 模型的 Core ML 轉換包。
transformers-to-coreml 的更新版本，一個基於 exporters 構建的無程式碼 Core ML 轉換工具。
一些轉換好的模型，例如 Llama 2 7B 或 Falcon 7B，可直接用於這些文字生成工具。

任務概覽

當我釋出推文展示 Falcon 或 Llama 2 在我的 Mac 上執行時，我收到了許多其他開發者的提問，他們想知道如何將這些模型轉換為 Core ML，因為他們也想在自己的應用中使用它們。轉換是一個關鍵步驟，但這只是拼圖的第一塊。我寫這些應用的真正原因是為了面對任何其他開發者都會遇到的問題，並找出我們可以在哪些方面提供幫助。在這篇文章的其餘部分，我們將介紹其中一些任務，並解釋我們在哪些方面有（或沒有）工具可以提供幫助。

轉換為 Core ML。我們將以 Llama 2 為例進行實際操作。
最佳化技術，使您的模型（和應用）執行更快並消耗盡可能少的記憶體。這是一個貫穿整個專案且沒有一蹴而就的解決方案的領域。
swift-transformers，我們用於幫助處理一些常見任務的新庫。
- 分詞器。分詞是將文字輸入轉換為模型處理的實際數字集（以及從生成的預測結果轉回文字）的方法。這比聽起來要複雜得多，因為有許多不同的選項和策略。
- 模型和 Hub 包裝器。如果我們想支援 Hub 上各種各樣的模型，就不能硬編碼模型設定。我們建立了一個簡單的 LanguageModel 抽象以及各種工具，用於從 Hub 下載模型和分詞器配置檔案。
- 生成演算法。語言模型被訓練來預測在一段文字序列後可能出現的下一個詞元的機率分佈。我們需要多次呼叫模型來生成文字輸出，並在每一步選擇一個詞元。決定下一步應該選擇哪個詞元有很多方法。
- 支援的模型。並非所有模型系列都（尚）受支援。
swift-chat。這是一個小型應用，簡單展示瞭如何在專案中使用 swift-transformers。
缺失部分/未來計劃。一些重要但尚未實現的功能，作為未來工作的方向。
資源。所有專案和工具的連結。

轉換為 Core ML

Core ML 是蘋果的原生機器學習框架，也是其使用的檔案格式的名稱。將模型從（例如）PyTorch 轉換為 Core ML 後，您就可以在 Swift 應用中使用它。Core ML 框架會自動選擇最適合執行模型的硬體：CPU、GPU 或稱為神經網路引擎（Neural Engine）的專用張量單元。根據您的系統特性和模型細節，也可能組合使用這些計算單元。

為了瞭解在實際中轉換模型是什麼樣的，我們將看看如何轉換最近釋出的 Llama 2 模型。這個過程有時可能很複雜，但我們提供了一些工具來提供幫助。這些工具並不總是有效，因為新模型層出不窮，我們需要不斷進行調整和修改。

我們推薦的方法是

使用 transformers-to-coreml 轉換 Space

這是一個基於 exporters（見下文）構建的自動化工具，它對您的模型要麼有效，要麼無效。它不需要編碼：輸入 Hub 模型識別符號，選擇您計劃使用模型的任務，然後點選應用。如果轉換成功，您可以將轉換後的 Core ML 權重推送到 Hub，就大功告成了！

您可以訪問該 Space 或直接在這裡使用

使用 exporters，一個基於 Apple 的 coremltools（見下文）構建的 Python 轉換包。

這個庫為您提供了更多配置轉換任務的選項。此外，它還允許您建立自己的轉換配置類，您可以用它來進行額外的控制或解決轉換問題。

使用 coremltools，Apple 的轉換包。

這是最底層的方法，因此提供了最大的控制權。對於某些模型（尤其是新模型），它仍然可能失敗，但您總是有機會深入原始碼，嘗試找出原因。

關於 Llama 2 的好訊息是，我們已經完成了前期工作，使用上述任何一種方法都可以成功轉換。壞訊息是，它在釋出時*轉換失敗*，我們不得不進行一些修復才能支援它。我們在附錄中簡要回顧了當時發生的情況，這樣您就可以瞭解當事情出錯時該怎麼做。

重要的經驗教訓

我跟蹤了一些最近模型（Llama 2、Falcon、StarCoder）的轉換過程，並將我學到的知識應用到了 exporters 和 transformers-to-coreml Space 中。以下是一些要點的總結

如果您必須使用 coremltools，請使用最新版本：7.0b1。儘管技術上是 beta 版，但我已經使用它幾周了，它真的很好：穩定，包含大量修復，支援 PyTorch 2，並有高階量化工具等新功能。
exporters 在轉換文字生成任務時不再對輸出應用 softmax。我們意識到這對於某些生成演算法是必要的。
exporters 現在預設對文字模型使用固定序列長度。Core ML 有一種方法可以指定“靈活形狀”，這樣您的輸入序列可以有任何長度，例如在 1 到 4096 個詞元之間。我們發現靈活輸入只能在 CPU 上執行，而不能在 GPU 或神經網路引擎上執行。更多調查即將展開！

我們將繼續向我們的工具中新增最佳實踐，這樣您就不必再次發現同樣的問題。

最佳化

如果轉換後的模型不能在您的目標硬體上快速執行並尊重系統資源，那麼轉換模型就毫無意義。本文提到的模型對於本地使用來說相當大，我們有意使用它們來挑戰當前技術的極限，並瞭解瓶頸所在。

我們確定了幾個關鍵的最佳化領域。它們對我們來說是一個非常重要的課題，也是當前和未來工作的主題。其中一些包括

快取先前生成的注意力鍵和值，就像 transformers 模型在 PyTorch 實現中所做的那樣。注意力分數的計算需要在迄今為止生成的整個序列上執行，但所有過去的鍵值對都已在之前的執行中計算過了。我們目前*沒有*為 Core ML 模型使用任何快取機制，但計劃這樣做！
使用離散形狀而不是小的固定序列長度。不使用靈活形狀的主要原因是它們與 GPU 或神經網路引擎不相容。次要原因是沒有快取，如上所述，生成會隨著序列長度的增長而變慢。使用一組離散的固定形狀，再加上快取鍵值對，應該可以實現更大的上下文大小和更自然的聊天體驗。
量化技術。我們已經在 Stable Diffusion 模型的背景下探索過它們，並對它們可能帶來的選擇感到非常興奮。例如，6 位調色盤化可以減小模型大小並且資源效率高。混合位量化是一種新技術，可以在對模型質量影響很小的情況下實現 4 位量化（平均）。我們也計劃在語言模型上研究這些主題！

對於生產應用，請考慮使用較小的模型進行迭代，尤其是在開發期間，然後應用最佳化技術來選擇您用例所能承受的最小模型。

`swift-transformers`

swift-transformers 是一個正在開發中的 Swift 包，旨在為 Swift 開發者提供類似 transformers 的 API。讓我們看看它有什麼以及缺少什麼。

分詞器

分詞解決了兩個互補的任務：將文字輸入適配為模型使用的張量格式，並將模型的結果轉換回文字。這個過程很微妙，例如

我們應該使用單詞、字元、字元組還是位元組？
我們應該如何處理小寫和大寫字母？我們是否應該處理這種差異？
我們應該移除重複的字元，比如空格嗎？或者它們是重要的？
我們如何處理不在模型詞彙表中的單詞？

有幾種通用的分詞演算法，以及許多不同的規範化和預處理步驟，這些對於有效使用模型至關重要。transformers 庫決定將所有這些操作抽象到同一個庫（tokenizers）中，並將這些決策表示為與模型一起儲存在 Hub 上的配置檔案。例如，這是 Llama 2 分詞器配置中描述*僅規範化步驟*的摘錄

  "normalizer": {
    "type": "Sequence",
    "normalizers": [
      {
        "type": "Prepend",
        "prepend": "▁"
      },
      {
        "type": "Replace",
        "pattern": {
          "String": " "
        },
        "content": "▁"
      }
    ]
  },

它的解讀是這樣的：規範化是按順序應用的一系列操作。首先，我們在輸入字串前 `Prepend`（前置）字元 `_`。然後我們將所有空格替換為 `_`。潛在的操作列表非常龐大，它們可以應用於正則表示式匹配，並且必須以非常特定的順序執行。tokenizers 庫中的程式碼處理了 Hub 中所有模型的這些細節。

相比之下，在其他領域（如 Swift 應用）中使用語言模型的專案，通常會採取將這些決策硬編碼到應用原始碼中的方式。對於少數幾個模型來說這沒問題，但這樣一來就很難用不同的模型替換現有模型，而且很容易出錯。

我們在 swift-transformers 中所做的，是在 Swift 中複製這些抽象，這樣我們只需要編寫一次，每個人都可以在他們的應用中使用它們。我們才剛剛開始，所以覆蓋範圍還很小。歡迎在倉庫中提交 issue 或貢獻您自己的程式碼！

具體來說，我們目前支援 BPE（位元組對編碼）分詞器，這是當今使用的三種主要系列之一。GPT 模型、Falcon 和 Llama 都使用這種方法。對 Unigram 和 WordPiece 分詞器的支援將在稍後推出。我們還沒有移植所有可能的規範化器、預分詞器和後處理器——只移植了我們在轉換 Llama 2、Falcon 和 GPT 模型時遇到的那些。

這是在 Swift 中使用 Tokenizers 模組的方法

import Tokenizers

func testTokenizer() async throws {
    let tokenizer = try await AutoTokenizer.from(pretrained: "pcuenq/Llama-2-7b-chat-coreml")
    let inputIds = tokenizer("Today she took a train to the West")
    assert(inputIds == [1, 20628, 1183, 3614, 263, 7945, 304, 278, 3122])
}

然而，您通常不需要自己對輸入文字進行分詞——Generation 程式碼會處理好這件事。

模型和 Hub 包裝器

如上所述，transformers 大量使用儲存在 Hub 上的配置檔案。我們準備了一個簡單的 Hub 模組來從 Hub 下載配置檔案，用於例項化分詞器和檢索有關模型的元資料。

關於模型，我們建立了一個簡單的 LanguageModel 型別作為 Core ML 模型的包裝器，專注於文字生成任務。透過使用協議，我們可以用相同的 API 查詢任何模型。

為了檢索您使用的模型的適當元資料，swift-transformers 依賴於一些自定義元資料欄位，這些欄位必須在轉換 Core ML 檔案時新增。swift-transformers 將使用這些資訊從 Hub 下載所有必要的配置檔案。以下是我們在 Xcode 的模型預覽中使用的欄位

exporters 和 transformers-to-coreml 會自動為您新增這些欄位。如果您手動使用 coremltools，請確保自己新增它們。

生成演算法

語言模型被訓練來預測在輸入序列後可能出現的下一個詞元的機率分佈。為了構成一個響應，我們需要多次呼叫模型，直到它產生一個特殊的*終止*詞元，或者我們達到我們期望的長度。決定下一個最佳詞元有很多方法。我們目前支援其中兩種

貪婪解碼（Greedy decoding）。這是最直接的演算法：選擇機率最高的詞元，將其附加到序列中，然後重複。對於相同的輸入序列，這將總是產生相同的結果。
top-k 取樣。選擇 `top-k`（其中 `k` 是一個引數）個最可能的詞元，然後使用像 `temperature` 這樣的引數從中隨機*取樣*，這將增加變異性，但代價是可能導致模型跑題並失去對內容的跟蹤。

其他方法如“核取樣”將在稍後推出。我們推薦閱讀這篇博文（最近更新），它對生成方法及其工作原理進行了出色的概述。像輔助生成這樣的複雜方法對於最佳化也可能非常有用！

支援的模型

到目前為止，我們已經用少數幾個模型測試了 swift-transformers，以驗證主要的設計決策。我們期待嘗試更多模型！

Llama 2。
Falcon。
StarCoder 模型，基於 GPT 架構的一個變體。
GPT 系列，包括 GPT2、distilgpt、GPT-NeoX、GPT-J。

`swift-chat`

swift-chat 是一個基於 swift-transformers 構建的簡單演示應用。其主要目的是展示如何在您的程式碼中使用 swift-transformers，但它也可以用作模型測試工具。

要使用它，請從 Hub 下載一個 Core ML 模型或建立您自己的模型，然後從 UI 中選擇它。所有相關的模型配置檔案將從 Hub 下載，使用元資料資訊來識別這是哪種模型型別。

首次載入新模型時，需要一些時間來準備。在此階段，CoreML 框架將編譯模型並根據您的機器規格和模型結構決定在哪些計算裝置上執行它。這些資訊會被快取並在未來的執行中重複使用。

該應用有意設計得簡單，以使其易於閱讀和簡潔。它也缺少一些功能，主要是由於當前模型上下文大小的限制。例如，它沒有任何“系統提示”的設定，這些提示對於指定您的語言模型的行為甚至其個性非常有用。

缺失部分/未來計劃

如前所述，我們才剛剛開始！我們接下來的優先事項包括

編碼器-解碼器模型，如 T5 和 Flan。
更多分詞器：支援 Unigram 和 WordPiece。
額外的生成演算法。
支援鍵值快取以進行最佳化。
使用離散序列形狀進行轉換。與鍵值快取相結合，這將允許更大的上下文。

讓我們知道您認為我們接下來應該做什麼，或者前往倉庫檢視 “Good First Issues” 來一試身手！

結論

我們介紹了一套工具，幫助 Swift 開發者將語言模型整合到他們的應用中。我迫不及待地想看看您用它們創造出什麼，並期待在社群的幫助下改進它們！請隨時與我聯絡 :)

附錄：以硬核方式轉換 Llama 2

除非您遇到過 Core ML 轉換問題並準備好與之戰鬥，否則您可以安全地忽略此部分 :)

根據我的經驗，PyTorch 模型使用 coremltools 轉換到 Core ML 失敗的兩個常見原因是

不支援的 PyTorch 操作或操作變體

PyTorch 有*大量*操作，所有這些操作都必須對映到一箇中間表示（MIL，即*模型中間語言*），然後該中間表示又被轉換為原生的 Core ML 指令。PyTorch 的操作集不是靜態的，所以新的操作也必須新增到 coremltools 中。此外，一些操作非常複雜，可以在其引數的各種奇特組合上工作。一個最近新增的、非常複雜的操作示例是*縮放點積注意力*，在 PyTorch 2 中引入。一個部分支援的操作示例是 einsum：並非所有可能的方程都能翻譯成 MIL。

邊緣情況和型別不匹配

即使對於支援的 PyTorch 操作，也很難確保轉換過程在所有可能的輸入和所有不同的輸入型別上都能正常工作。請記住，單個 PyTorch 操作可以有多個針對不同裝置（cpu、CUDA）、輸入型別（整數、浮點數）或精度（float16、float32）的後端實現。所有組合的乘積是驚人的，有時模型使用 PyTorch 程式碼的方式會觸發一個可能沒有被考慮或測試過的轉換路徑。

這就是我第一次嘗試使用 coremltools 轉換 Llama 2 時發生的情況

透過比較不同版本的 transformers，我發現問題是在引入這行程式碼時開始出現的。這是最近一次 `transformers` 重構的一部分，旨在更好地處理*所有*使用因果掩碼的模型中的因果掩碼，所以這對其他模型來說也是一個大問題，不僅僅是 Llama。

錯誤截圖告訴我們，在填充掩碼張量時存在型別不匹配。問題來自於行中的 `0`：它被解釋為一個 `int`，但要填充的張量包含 `float`，而使用不同型別被轉換過程拒絕了。在這種特殊情況下，我為 coremltools 提出了一個補丁，但幸運的是這種情況很少需要。在許多情況下，您可以修補您的程式碼（在 `transformers` 的本地副本中使用 `0.0` 就可以了），或者建立一個“特殊操作”來處理異常情況。我們的 exporters 庫對自定義的特殊操作有很好的支援。請參閱這個例子，瞭解一個缺失的 `einsum` 方程，或者這個例子，瞭解一個使 `StarCoder` 模型在新版 `coremltools` 釋出前能正常工作的變通方法。

幸運的是，coremltools 對新操作的覆蓋範圍很好，團隊反應也非常快。