WWDC 24: 使用 Core ML 執行 Mistral 7B

WWDC’ 24 是蘋果正式釋出 Apple Intelligence 並重申其對高效、私密、裝置端 AI 承諾的時刻。在主題演講及隨後的分會場中，他們展示了 Apple Intelligence，它驅動了大量 AI 增強功能，展示了 AI 在日常任務中的實際應用。這些不是*為了 AI 而 AI* 的華麗演示，而是能節省時間、恰當（且有趣！）的助手，它們深度集成於應用程式和作業系統中，併為開發者提供了多種方式將這些功能融入他們自己的應用中。

Apple Intelligence 的各項功能之所以能如此出色，得益於其垂直整合的軟體棧，它能夠最大限度地發揮 Apple Silicon 的能力。蘋果還為開發者提供了一個在裝置上執行模型的平臺，即 Core ML。這個軟體棧允許你在 Apple Silicon 硬體上所有三種計算單元（CPU、GPU 和神經網路引擎）上執行機器學習模型。

在本篇博文中，我們將探討一些 Core ML 最出色的新功能，以復現蘋果在 WWDC'24 使用 Core ML 在裝置端部署機器學習和 AI 模型分會場中展示的 Mistral 7B 示例。在該示例中，他們使用了一個 swift-transformers 的 fork 版本，在 Mac 上運行了一個頂尖的 LLM。這是一個擁有超過 70 億引數的高質量模型，它挑戰了當今消費級硬體的極限。你也可以檢視 WWDC'24 的將你的機器學習和 AI 模型引入 Apple silicon 分會場，其中展示了 Mistral 7B 轉換過程的一部分。

讓我們看看需要採取哪些步驟來儘可能高效地執行它，並學習 iOS 18 和 macOS Sequoia 中可用的新工具。

這是我們今天要構建的內容

太長不看版

讀完這篇博文，你將瞭解最新 macOS 版本帶來的所有新特性，併成功地在你的 Mac 上用不到 4GB 的記憶體執行一個 70 億引數的模型。

第一步：克隆 swift-transformers 倉庫的 preview 分支：git clone -b preview https://github.com/huggingface/swift-transformers 第二步：從 此 Hugging Face 倉庫 下載轉換後的 Core ML 模型 第三步：使用 Swift 執行推理：swift run transformers "2024年8月在巴黎遊覽的最佳推薦地點是：" --max-length 200 Mistral7B-CoreML/StatefulMistralInstructInt4.mlpackage

WWDC' 24 最亮眼的 Core ML 新功能

以下是我們在 Mac 上執行 Mistral 7B 將用到的一些來自 WWDC' 24 的最具影響力的 Core ML 功能。

Swift Tensor

我們首先要強調的功能是一個全新的 Swift 型別，用於處理 ML 張量。這些是每個 ML 框架都使用的多維資料結構。從事 ML 的 Python 開發者對 numpy 陣列或 torch 張量很熟悉，它們為輕鬆操作這些大型多維矩陣提供了便捷的高階介面。新的 MLTensor 型別提供了一個類似於 Python 框架中可用的高階抽象，極大地簡化了在 Swift 中處理張量資料的工作。

Core ML 之前已經有多維資料型別，如 MLMultiArray 和 MLShapedArray。然而，它們主要用於資料儲存和簡單操作，比如將你的資料打包並作為輸入傳送給 Core ML 模型，或者從 Core ML 模型中解包結果。但是，使用這些 API 來*操縱*張量資料是很困難的。它們只提供了一些基本操作，你可能需要透過訪問底層儲存作為指向數字資料的不透明指標來編寫自己的操作。這既耗時又容易出錯。

新的 Tensor 型別提供了一個高階抽象，模仿了 Python 框架中可用的抽象，極大地簡化了在 Swift 中處理張量資料的工作。考慮一個像我們想要移植到 Core ML 的語言模型。語言模型接收一個 token 輸入序列，並輸出對詞彙表中所有 token 機率的估計，這意味著機率高的 token 很有可能是輸入的合理延續。應用程式的工作是根據這些機率選擇最佳的下一個 token 來追加到序列中。Tensor 型別使得處理這些操作變得容易，無需自定義程式碼。

當我們釋出 swift-transformers 時，我們編寫了大量程式碼（後來由社群擴充套件，謝謝！❤️）來幫助進行輸入準備（將單詞轉換為 token）和輸出後處理。例如，可以檢視我們使用 Accelerate 實現的 softmax 操作。所有這些在使用 MLTensor 時都可以移除，因為 softmax 是開箱即用的！

有狀態緩衝區 (Stateful Buffers)

在 WWDC’ 24 之前，Core ML 模型本質上是一個純粹的無狀態函式，你提供輸入並返回一些輸出。然而，有時你需要保持一個依賴於先前計算的狀態。維護狀態的函數語言程式設計方法是增加一個額外的輸入/輸出對。因此，模型根據你的輸入和狀態計算輸出和新狀態。這種方法沒有錯，事實上，像 JAX 這樣的高效能框架就是這樣工作的。

然而，這種方法存在實際限制：每次呼叫模型時，都需要將有狀態的資料作為輸入傳送給模型，並作為輸出取回。如果有狀態的資料很大，這種來回傳輸會增加開銷並降低速度。這對於 LLM 尤其重要，因為你需要執行多次迭代來生成一個序列。效能瓶頸通常是你的計算機記憶體頻寬（即，你將資料移至 GPU 及返回的速度有多快）。有狀態模型透過為狀態資料預留一塊記憶體並將其保留在 GPU 上來解決這個問題，這樣你就不必每次使用模型時都發送和接收它了。

有狀態緩衝區是在 這次 WWDC' 24 分會場 中透過一個易於理解但不代表像 LLM 這樣大型模型的實際應用的玩具示例來介紹的。對於基於 Transformer 的模型，一個 LLM 效能技巧是鍵值快取（稱為 kv-caching）。如下圖所示，它透過快取先前步驟中執行的操作結果來避免關鍵的注意力模組中昂貴的矩陣乘法。我們不會深入細節，但關鍵點是：kv-cache 顯著提高了效能，並且它需要一大塊記憶體，這是使用有狀態緩衝區的完美候選者。這裡有一份關於有狀態模型的 coremltools 使用者指南更新。

新的量化技術

在 WWDC 23 中，我們探討了一種非常酷的技術，稱為調色盤化 (palletization)，並展示了它如何幫助將文生圖模型，例如 Stable Diffusion，帶到 Mac 和 iPhone 上。

雖然這些技術可以讓你大幅減小模型尺寸，但如果壓縮得太過，對質量的影響是巨大的。更大的模型受此影響更嚴重，因為權重資料的動態範圍很廣。建立一個能夠捕捉所有可能值的小型查詢表（LUT）變得越來越困難。WWDC 24 中引入的解決方案是每次關注資料的一小部分，併為同一張量的不同區域建立多個查詢表。

這些方法（分塊量化）使我們能夠將模型壓縮到低至 4 位精度。我們不再使用 4 位元組（一個 float32 數字的大小）來表示每個模型引數，而是可以只用半個位元組（一個 nibble）。這使得模型大小減少了 8 倍（減去一些用於分塊量化表的開銷），或者與 float16 精度相比小了 4 倍。

多函式支援

雖然這個例子我們不會用到這個功能，但我們想在這裡提一下，因為它是在 WWDC 24 上推出的，我們將在一些未來的工作中展示它。多函式支援基本上允許你將 LoRA 介面卡打包到生成模型中，以便為不同的任務使用同一個模型（只需一小組額外的引數，稱為介面卡）。LoRA 是社群首選的大模型微調技術。例如，在擴散模型中，你可以使用 LoRA 生成不同風格的影像，如照片般逼真或卡通風格。我們相信 LoRA 是驅動蘋果 Genmoji 實現的解決方案的一部分。對於語言模型，LoRA 介面卡可用於將通用 LLM 調整到特定任務或領域。

要了解更多關於 LoRA 的資訊，你可以檢視這篇文章。

要了解更多關於多函式支援的資訊，你可以檢視蘋果 coremltools 使用者指南這裡。

將 Mistral 7B 轉換為 Core ML

高效執行大型語言模型最重要的單一元件是 kv-cache。如上所述，這是 WWDC' 24 釋出的新有狀態模型特性的絕佳候選者。transformers 庫中的模型已經使用了高效的注意力實現，這些實現嚴重依賴於 kv-caching。然而，預設實現是為 Nvidia GPU 最佳化的，而這種硬體與 Apple Silicon 有著不同的約束。對於 Core ML，我們需要預先分配完整的快取緩衝區，並確保每次呼叫模型時，我們都能就地更新緩衝區。這避免了低效的記憶體分配和張量拼接，同時也是 Core ML 有狀態緩衝區的要求。

為了實現這個目標，我們必須使用一種考慮了這些因素的不同注意力實現。這需要修改 Mistral 架構的 transformers 建模程式碼，這是在這段程式碼片段中完成的。

注意：如果你想跟著做並復現轉換過程（或轉換另一個基於 Mistral 的模型，例如不同的微調版本），你可以使用這個指令碼來執行所有的轉換步驟。

追蹤與轉換

第一步是載入模型。我們將使用帶有就地快取 (in-place cache) 方法的補丁實現。

MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
torch_model = StatefulMistralForCausalLM(MODEL_ID)
torch_model.eval()

在執行 Core ML 轉換之前，我們需要用示例輸入來追蹤模型。這個過程會記錄在這些輸入上執行的張量操作，追蹤到的計算圖將在轉換期間被翻譯成 Core ML 操作。我們使用樣本輸入來追蹤模型；我們不需要真實資料。

input_ids = torch.zeros((1, 2), dtype=torch.int32)
causal_mask = torch.zeros((1, 1, 2, 5), dtype=torch.float32)

traced_model = torch.jit.trace(torch_model, [input_ids, causal_mask])

語言模型的輸入是一個長度可變的 token 序列。我們將允許輸入從單個 token 增長到最大上下文長度 2048。我們可以使用 coremltools 的範圍維度來指定這些界限。

query_length = ct.RangeDim(lower_bound=1, upper_bound=2048, default=1)
end_step_dim = ct.RangeDim(lower_bound=1, upper_bound=2048, default=1)

inputs = [
    ct.TensorType(shape=(1, query_length), dtype=np.int32, name="inputIds"),
    ct.TensorType(shape=(1, 1, query_length, end_step_dim), dtype=np.float16, name="causalMask"),
]

outputs = [ct.TensorType(dtype=np.float16, name="logits")]

除了序列 token（在上面的例子中稱為 inputIds），還有另一個輸入叫做 causalMask，它指定了模型需要關注哪些 token。這主要用於在使用批處理同時生成多個序列時。可以檢視這些輸入在一個示例執行器中的用法。

在這種情況下，一個批次內的所有輸入序列必須具有相同的長度，所以我們使用填充 token 和因果掩碼來告訴模型，填充 token 不應被視為輸入。

狀態準備

PyTorch 建模程式碼使用 keyCache 和 valueCache 作為快取緩衝區的名稱來儲存 kv-cache。這些塊是為最大上下文長度（2048）分配的。我們使用 coremltools 新的 StateType 來指定這些塊在轉換期間必須轉換為有狀態的 Core ML 緩衝區。

# Specify kv-cache states by using `StateType`.

states = [
    ct.StateType(
        wrapped_type=ct.TensorType(shape=torch_model.kv_cache_shape, dtype=np.float16),
        name="keyCache",
    ),
    ct.StateType(
        wrapped_type=ct.TensorType(shape=torch_model.kv_cache_shape, dtype=np.float16),
        name="valueCache",
    ),
]

Core ML 轉換

要將模型轉換為 Core ML，我們需要指定輸入和輸出型別以及狀態。轉換後的模型將使用 float16 精度，因為這是我們為輸入資料指定的。我們還需要將最低部署目標指定為 iOS18，因為這些功能是在該版本中提供的。（我們也可以使用 macOS15，它指向相同的轉換目標。）

mlmodel_fp16 = ct.convert(
    traced_model,
    inputs=inputs,
    states=states,
    outputs=outputs,
    minimum_deployment_target=ct.target.iOS18,
    skip_model_load=True,
)

模型壓縮

使用上述新的分塊量化策略，我們採用塊大小為 32 的 4 位線性量化。這將大大減小模型大小並使模型執行更快。儘管計算仍將以 float16 格式進行，但權重以 4 位模式傳輸並在執行時動態解壓，這比傳輸大量的 16 位權重更高效。

量化引數配置如下

op_config = ct.optimize.coreml.OpLinearQuantizerConfig(
    mode="linear_symmetric",
    dtype="int4",
    granularity="per_block",
    block_size=32,
)
config = ct.optimize.coreml.OptimizationConfig(global_config=op_config)

讓我們用這個配置來量化模型。下面這行程式碼將需要幾分鐘才能執行完畢。

mlmodel_int4 = ct.optimize.coreml.linear_quantize_weights(mlmodel_fp16, config=config)

mlmodel_int4.save("StatefulMistral7BInstructInt4.mlpackage")

在轉換和量化完成後，還有最後一步。我們需要包含一個額外的元資料，指明我們使用的模型識別符號（mistralai/Mistral-7B-Instruct-v0.3）。Swift 程式碼將使用這個識別符號從 Hub 下載 tokenizer 檔案。Tokenization 是將文字資料轉換為模型使用的數值表示的過程，每個模型都不同。

mlmodel_int4._spec.description.metadata.userDefined.update({
    "co.huggingface.exporters.name": MODEL_ID
})

生成的模型是一個約 3.8G 的 mlpackage，而 float16 轉換會產生 14G 的檔案。你可以在 Hub 上找到它。

使用 Swift 執行 Mistral 7B

如果你按照上述步驟操作或從 Hub 下載了模型，你可以使用 swift-transformers 的 preview 分支在本地執行它。蘋果工程師向該專案貢獻了程式碼，包括以下重要功能

• 完整的 Tensor 支援，極大地簡化了預處理和後處理任務，並允許我們刪除許多低階、混亂且脆弱的程式碼。

• 支援 Stateful API 的 Swift 對應版本。

由於採用這些功能是破壞性改動，且需要 iOS 18 或 macOS 15，我們暫時將它們保留在 preview 分支中。

要從命令列執行模型，請先從 GitHub 倉庫克隆 preview 分支

    git clone -b preview https://github.com/huggingface/swift-transformers

然後執行 CLI 測試模型

#to run in release mode, pass -c release
swift run transformers "Best recommendations for a place to visit in Paris in August 2024:" --max-length 128 Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage

為了方便測試，你還可以使用 swift-chat，這是一個我們編寫的簡單應用，用於展示如何將 swift-transformers 包整合進去。你也必須使用 preview 分支。本文開頭展示了執行轉換後的 Mistral 模型的 swift-chat 示例。

使用 Python 執行 Mistral 7B

對於更熟悉 Python 的朋友來說，這也同樣簡單！

python3 generate.py Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage --prompt "Best recommendations for a place to visit in Paris in August 2024:"

coremltools 使得用 Python 執行 Core ML 模型變得同樣容易。

下一步是什麼？

我們對今年 Core ML 和 coremltools 的進展感到非常興奮，並期待看到大量第三方應用利用 ML 模型來解決人們實際需要的任務。在我們這邊，我們致力於讓這一切儘可能簡單，以便開發者可以專注於創造酷炫的應用。我們正在籌劃幾件事

• 這裡介紹的模型更新非常適合 Mac 電腦上的 GPU。Core ML 可以使用神經網路引擎，這在 iPhone 上尤其高效。要從神經網路引擎中獲得最佳效能，需要一些額外的適配，我們計劃在一些示例模型上進行這些工作。這項工作將基於這篇 2022 年（至今仍然非常相關）的蘋果文章中討論的經驗。我們不會在 iPhone 上執行 Mistral 7B，但有幾個較小的模型，如蘋果的 OpenELM 或 DCLM，是探索的絕佳候選者！

• 這裡展示的程式碼是高度實驗性的。隨著夏天的進行，我們計劃採用這些方法並將其整合到 exporters 中，這是一個旨在將 transformers 模型轉換為 Core ML 的 Python 工具。希望你很快就能非常輕鬆地轉換許多有趣的模型架構。

• 我們將繼續在 swift-transformers 的 preview 分支上工作，以整合新功能或 API 變更。如果你感興趣，請關注它！

你能如何幫助我們？

蘋果在 WWDC 上釋出的工具幫助我們實現了讓 AI 對所有人來說都簡單易用的長期目標，我們很樂意看到你將它們帶向何方。我們展示的例子是實驗性的，但你可以用它來將任何 Mistral 微調模型轉換為 Core ML——如果你這麼做了，請告訴我們！如果你想嘗試其他模型架構，請隨時向 swift-transformers 的 preview 分支提出問題或 PR——我們會盡力幫助你開始！

現在是運用你的創造力來解決你感興趣的問題的最佳時機！去嘗試，去享受，並告訴我們如何能幫助你。