Core ML で Mistral 7B を実行する

Hugging Face は、WWDC 24 で導入された新しい Core ML 機能を活用して、Mac ハードウェア上で Mistral 7B モデルを実行するワークフローを詳しく紹介しています。ステートフルバッファ、新しい Swift Tensor 型、そして高度な量子化を組み合わせることで、70 億パラメータのモデルを 4GB 未満のメモリで実行できます。

WWDC 24 からの新しい Core ML 機能

Apple の最新の Core ML アップデートは、オンデバイスの大規模言語モデル（LLM）向けにいくつかの重要な最適化を提供します。

Swift Tensor (MLTensor)

新しい MLTensor 型は、Swift における多次元データ構造の高レベル抽象を提供し、numpy 配列や torch テンソルの機能を模倣します。これにより、以前は不透明ポインタとして基底ストレージにアクセスする必要があった MLMultiArray や MLShapedArray の低レベル操作が不要になります。MLTensor には softmax などの組み込み操作が含まれており、言語モデルの前処理および後処理パイプラインを簡素化します。

ステートフルバッファ

従来、Core ML モデルはステートレス関数として動作していました。ステートフルバッファは、GPU 上にメモリブロックを確保してイテレーション間の状態を保持できるようにします。LLM では、キー・バリュー（KV）キャッシュを実装するために使用されます。KV キャッシュを GPU 上に保持することで、各トークン生成ステップで CPU と GPU 間で大きなステートフルデータを送受信するオーバーヘッドを回避し、メモリ帯域幅のボトルネックを削減してパフォーマンスを大幅に向上させます。

ブロック単位の量子化

品質を大幅に損なうことなくモデルサイズを削減するために、Core ML はブロック単位の量子化をサポートするようになりました。この手法は、同一テンソルの異なる領域ごとに複数のルックアップテーブル（LUT）を作成し、全体を単一のテーブルで表すのではありません。これにより、モデルを 4 ビット精度に圧縮でき、float32 に比べてサイズが 8 倍、float16 に比べて 4 倍削減されます。

マルチファンクションサポート

マルチファンクションサポートにより、開発者は LoRA（Low-Rank Adaptation）アダプタを生成モデルに組み込むことができます。これにより、追加パラメータ（アダプタ）の小さなセットを入れ替えるだけで、単一のベースモデルをさまざまなタスクやスタイルに利用でき、モデル全体を再ロードする必要がなくなります。

Mistral 7B を Core ML に変換する

Mistral 7B を効率的に実行するには、KV キャッシュバッファ全体を事前に割り当て、インプレースで更新するカスタムアテンション実装が必要で、これにより Core ML のステートフルバッファ要件を満たします。

トレースと変換プロセス

  1. Tracing: モデルは StatefulMistralForCausalLM のパッチ適用実装を使用してロードされ、torch.jit.trace を用いてサンプル入力でトレースされます。
  2. Input Definition: coremltools を通じて範囲次元が使用され、入力シーケンスが単一トークンから最大コンテキスト長 2048 まで拡張できるようにします。
  3. State Preparation: keyCache と valueCache バッファは ct.StateType を使用して定義され、ステートフルな Core ML バッファに変換されることが保証されます。
  4. Conversion: モデルは iOS 18 または macOS 15 を最低デプロイメントターゲットとして Core ML に変換され、新しいステートフル機能にアクセスできるようになります。

モデル圧縮

OpLinearQuantizerConfig を使用して、モデルはブロックサイズ 32 の 4 ビット線形対称量子化で圧縮されます。これにより、最終的な mlpackage のサイズは約 14GB（float16 時）から約 3.8GB に削減されます。

実行と実装

ユーザーは swift-transformers リポジトリの preview ブランチを使用して、変換された Mistral 7B モデルを実行できます。

Swift での実行

preview ブランチは完全な MLTensor サポートと、ステートフル API の Swift 対応版を組み込んでいます。推論はコマンドラインから実行できます。

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

Python での実行

推論は coremltools を使用した Python でもサポートされています。

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

今後のロードマップ

Hugging Face は、これらの実験的手法を exporters（transformers モデルを Core ML に変換する Python ツール）に統合することを計画しています。さらに、現在の Mistral 7B 実装が主に Mac GPU 向けに最適化されているため、OpenELM や DCLM などの小型モデルを検討し、iPhone の Apple Neural Engine（ANE）向けのパフォーマンス最適化も目指しています。

Sources