ABEJA でエンジニアをしている石川です。これは ABEJA アドベントカレンダー 2024 の 18 日目の記事です。

macOS/iOS で機械学習モデルを動かすにはいくつかの方法がありますが、Apple シリコンの能力を十分に引き出すためには CoreML を使うのが最適です。

Python 向け機械学習フレームワークである PyTorch も MPS バックエンドによって、Apple シリコンの GPU を利用することはできます。しかし、Apple の NPU (Neural Processing Unit: 機械学習やディープラーニングなどの計算処理を高速化するために設計された専用プロセッサ) である Neural Engine ¹ を使えるのは CoreML だけです。

CoreML で機械学習モデルを動かす

CoreML で機械学習モデルを動かすためには、PyTorch や TensorFlow で作られたモデルを Core ML Tools で変換する必要があります。最近はインターネット上に情報も増え、HuggingFace の exporters といったツールも出てきたため、昔と比べると間口が広くなってきました。

また、macOS/iOS の機械学習サポートも進化をつづけており、最新の macOS15/iOS18 では LLM を動かすために効果的なアップデートがいくつも入っています。一例を挙げると:

Scaled Dot Product Attention (SDPA) オペレーションの追加 ²
Stateful Models によってモデル内部で in-place の更新が可能に ³
Int4 量子化

また、Swift に MLTensor が追加され、Python の numpy 配列や torch Tensor のように高度な操作が可能になりました。⁴

しかし、言語モデルを CoreML 向けに変換しただけで、テキスト生成ができるわけではありません。CoreML でテキスト生成を行うためには、モデル変換だけでなく、実際の推論プロセスを工夫する必要があります。特に以下のような課題を解決する必要があります。

トークン化の処理
テキスト生成では、入力のテキストをトークン（単語や部分単語などの単位）に変換し、モデルに与える必要があります。CoreML や Swift の標準ライブラリはトークン化処理をサポートしていません。

モデルの再帰的呼び出し
テキスト生成では、モデルから出力されたトークンを次の入力として再びモデルに与える再帰的な処理が必要です。CoreML のモデルは一度の推論で固定長の入力に対する出力を返すのが基本ですが、この制約を超えるためにはループ構造をコード側で組む必要があります。

サンプリングやビームサーチなどの実装
言語モデルは、テキスト生成の際に次に来るべき単語（トークン）を予測することができます。しかし、モデルが出力するのは「各トークンが選ばれる確率」を示しているにすぎません。そのため、予測された確率に基づいてどのトークンを選ぶかを決める方法と実装が必要です。⁵

`swift-transformers` を試す

上記の推論プロセスは、Python であれば HuggingFace の transformers ライブラリが提供してくれますが、実は Swift 向けにも swift-transformers というライブラリが提供されています。⁶

Mistral 7B モデルを動かす

実際に動かしてみましょう。今回は swift-transformers の macOS15/iOS18 対応版 ⁷ で mistralai/Mistral-7B-Instruct-v0.3 を動かします。なお、試した環境は以下の通りです。

MacBook Pro (13 インチ、M2、2022)
メモリ 16 GB
macOS 15.1.1

swift-transformers リポジトリを clone して preview ブランチにします。

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

CoreML 向けに変換済みの Mistral 7B モデルを apple/mistral-coreml から clone しておきます。

$ git clone https://huggingface.co/apple/mistral-coreml

さらに前準備として、本家 mistralai/Mistral-7B-Instruct-v0.3 にアクセスできるようにする必要があります。

mistralai/Mistral-7B-Instruct-v0.3 でモデルの利用規約に同意しておく
リポジトリへのアクセス権を付与したアクセストークンを発行する ⁸
トークンを HUGGING_FACE_HUB_TOKEN 環境変数に設定

$ export HUGGING_FACE_HUB_TOKEN=hf_...

次のコマンドで実行できます。

$ swift run transformers "What is Generative AI" --max-length 200 ../mistral-coreml/StatefulMistral7BInstructInt4.mlpackage

以下のような出力が得られました。

Generative AI is a type of artificial intelligence that creates new content, such as images, text, or music, by learning patterns from existing data. It generates new data that is similar in style and structure to the training data, but not an exact copy.

Generative AI models are trained on large datasets and use a process called "generative modeling" to learn the underlying distribution of the data. This allows the model to generate new data that is statistically similar to the training data, but with variations and new combinations.

Some examples of generative AI include:

1. Text generation: Generative AI can be used to write articles, stories, or even code. It can generate text that is similar in style and content to existing text, but with new ideas and perspectives.
2. Image generation: Generative AI can create new images that are similar in style and content to existing images, but with new variations and combinations. This can be
14.97 tokens/s, prompt pre-filling time: 1.42s, total time: 14.71s

最新の CoreML アップデートの詳細は HuggingFace のブログや Apple のブログを読むことで概要や使い方をつかむことができます。

`swift-transformers` で推論を実装する

ここで記事を終えてもいいのですが、今回はせっかくなので、LLM（大規模言語モデル）の祖先とも言える ⁹ GPT-2 ¹⁰ を CoreML モデルに変換して、swift-transformers のプレビュー版で動かしてみます。

Python で動かしてみる

Swift で動かす前に、Python でどれくらいのパフォーマンスが出るのか見てみましょう。

model = GPT2LMHeadModel.from_pretrained("gpt2")

このスクリプトでは、transformers でテキストを逐次出力するための TextStreamer を拡張して、TTFT (Time To First Token) と TPS (Tokens Per Second) を計測しています。また、計測用のスクリプトでは、計測前に 3 回暖機運転を行い、その後 5 回実行した平均値を出すようにしています。

このスクリプトで、Python transformers による推論を計測した結果が以下の通りです。

poetry run python ./inference-metrics.py --warm 3 -n 5 -- poetry run python ./gpt-2/baseline/python-generate.py --max-length=106 "What is generative AI?"
...

Average Metrics:
[Prompt]  => 6 tokens, latency (TTFT): 0.29 ms
[Extend]  => 100 tokens, throughput: 44.93 tokens/s

毎秒 10 トークンを生成できれば、平均的な人が文字を読む速度よりも高速なので、44.93 tokens/s はかなり優秀ですね。

では、Python の transformers による実装では、Neural Engine が使われていないことを確認しましょう。powermetrics コマンドを使うと、CPU、GPU、ANE（Apple Neural Engine）の推定消費電力を表示することができます。

$ sudo powermetrics --samplers cpu_power,gpu_power,ane_power -i 1000
...
CPU Power: 11184 mW
GPU Power: 1251 mW
ANE Power: 0 mW
Combined Power (CPU + GPU + ANE): 12435 mW
...

Python 実装を動かしている間に、powermetrics で CPU/GPU/ANE の推定消費電力を収集した結果が次のグラフになります。 ¹¹

見ての通り、GPU は多少使われていますが、Neural Engine は一切使われていません。

さて、これを CoreML で動かせばどれだけ速くなるでしょうか？

CoreML モデルに変換

では、GPT-2 モデルを CoreML モデルに変換していきましょう。

Core ML Tools 8.0 以降では、PyTorch モデルを Core ML に変換する際に、torch.export という新しい方法も利用できるようになりました。 torch.export は、TorchScript よりも柔軟性が高く、将来的に TorchScript に取って代わる可能性を秘めています。しかしながら、torch.export は現時点ではまだ実験段階の機能です。そのため、今回は従来の TorchScript を用いた変換方法を紹介します。¹²

まず、入出力とモデルに渡すオプションを揃えるために、GPT2LMHeadModel のサブクラスを用意します。

class BaselineGPT2LMHeadModel(GPT2LMHeadModel):
    """Baseline LlamaForCausalLM model without key/value caching."""

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.LongTensor,
    ) -> torch.Tensor:
        out = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=False,
        )
        return out.logits

このサブクラスでは、親クラスのGPT2LMHeadModelのforwardメソッドをオーバーライドしています。主な変更点は以下の 2 点です。

use_cache=Falseを指定: CoreML はデフォルトではステートレスのため、use_cache をFalse に設定しています
out.logitsを返す: モデルの出力をシンプルにし、Core ML での取り扱いを容易にするため、logitsのみを返すようにしています

次に、バッチサイズを 1、コンテキスト長を 1024 とした入力サンプルを作成します。この設定では、入力の shape は (1, 1024) になります。input_ids と attention_mask には、今回はゼロ値のテンソルを使用します。

input_shape = (1, 1024)

example_inputs: tuple[torch.Tensor, torch.Tensor] = (
    torch.zeros(input_shape, dtype=torch.long),  # input_ids
    torch.zeros(input_shape, dtype=torch.long),  # attention_mask
)

最後に、作成した入力サンプルとモデルを使って、torch.jit.trace で TorchScript モデルを作成します。

traced_model: torch.jit.ScriptModule = torch.jit.trace(
    torch_model,  # 変換対象のモデル
    example_inputs=example_inputs,  # 入力サンプル
)

torch.jit.trace は、モデルに具体的な入力値を与えて実行し、その実行トレースを記録することで TorchScript モデルを作成する関数です。

これで TorchScript モデルへの変換が完了しました。次のステップでは、この TorchScript モデルを Core ML モデルに変換していきます。

inputs: list[ct.TensorType] = [
    ct.TensorType(shape=input_shape, dtype=np.int32, name="inputIds"),
    ct.TensorType(shape=input_shape, dtype=np.int32, name="attentionMask"),
]
outputs: list[ct.TensorType] = [ct.TensorType(dtype=np.float16, name="logits")]

CoreML はデフォルトで Float16 の精度で出力されるので、出力の型は float16 にしています。

mlmodel: ct.models.MLModel = ct.convert(
    traced_model,
    inputs=inputs,
    outputs=outputs,
    minimum_deployment_target=ct.target.iOS18,
    skip_model_load=True,
)

SDPA サポートなどの新しい機能を使うためには minimum_deployment_target に ct.target.iOS18 を指定します。¹³

最後に、変換したモデルを保存します。

mlmodel.save("models/gpt2-baseline.mlpackage")

変換用スクリプトはこちらです。いっぱい警告メッセージが出ますが無視して構いません。¹⁴

poetry run python ./gpt-2/baseline/export-model.py --context-size 1024 --minimum-deployment-target iOS18
/Users/takanori.ishikawa/Developer/Workspace/coreml-llm/.venv/lib/python3.12/site-packages/transformers/modeling_utils.py:4713: FutureWarning: `_is_quantized_training_enabled` is going to be deprecated in transformers 4.39.0. Please use `model.hf_quantizer.is_trainable` instead
  warnings.warn(
/Users/takanori.ishikawa/Developer/Workspace/coreml-llm/.venv/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py:114: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
/Users/takanori.ishikawa/Developer/Workspace/coreml-llm/.venv/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py:162: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if past_key_values_length > 0:
Converting PyTorch Frontend ==> MIL Ops:   0%|                                                                                                      | 0/1034 [00:00<?, ? ops/s]Saving value type of int64 into a builtin type of int32, might lose precision!
Saving value type of int64 into a builtin type of int32, might lose precision!
Converting PyTorch Frontend ==> MIL Ops: 100%|████████████████████████████████████████████████████████████████████████████████████████▉| 1033/1034 [00:00<00:00, 5650.27 ops/s]
Running MIL frontend_pytorch pipeline: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 126.76 passes/s]
Running MIL default pipeline:  63%|████████████████████████████████████████████████████████████████▊                                      | 56/89 [00:00<00:00, 66.98 passes/s]/Users/takanori.ishikawa/Developer/Workspace/coreml-llm/.venv/lib/python3.12/site-packages/coremltools/converters/mil/mil/ops/defs/iOS15/elementwise_unary.py:889: RuntimeWarning: overflow encountered in cast
  return input_var.val.astype(dtype=string_to_nptype(dtype_val))
/Users/takanori.ishikawa/Developer/Workspace/coreml-llm/.venv/lib/python3.12/site-packages/coremltools/converters/mil/mil/passes/defs/optimize_repeat_ops.py:433: RuntimeWarning: overflow encountered in cast
  max(cur_range.low, tmp_range.low), min(cur_range.high, tmp_range.high)
Running MIL default pipeline: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 89/89 [00:03<00:00, 22.35 passes/s]
Running MIL backend_mlprogram pipeline: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 12/12 [00:00<00:00, 112.23 passes/s]

念のため、出来上がったファイルを確認しておきましょう。

ls -lah models/gpt2.mlpackage/Data/com.apple.CoreML/weights/
total 483688
drwx------  3 takanori.ishikawa  staff    96B Dec 11 07:36 .
drwxr-xr-x  4 takanori.ishikawa  staff   128B Dec 11 07:36 ..
-rw-r--r--  1 takanori.ishikawa  staff   236M Dec 11 07:36 weight.bin

CoreML が出力する .mlpackage は macOS のパッケージなので、フォルダになっています。その中にある重みファイルの容量は 236MB。124M パラメータのモデルを Float16 精度で変換したので規模感は合っていそうです。

Xcode でも開いて情報を確認します。

Availability が iOS 18.0+ / macOS 15.0+ になっている
Operation に iOS18.scaledDotProductAttention が含まれている

ことが確認できました。

Swift で動かす

いよいよ Swift で動かしていきましょう。変換した Core ML モデルを読み込んで、推論プロセスを実行するコマンドラインツールを swift-transformers のプレビュー版を使って実装していきましょう。

ソースコード全体はこのフォルダにあるので、ここでは要点だけを解説していきます。¹⁵

CoreML モデルの読み込み

まずは、.mlpackage ファイルをMLModel として読み込める .mlmodelc ファイルにコンパイルしておきます。 .mlmodelc は、Core ML モデルをより効率的に実行するために最適化された形式です。

$ xcrun coremlcompiler compile models/gpt2-baseline.mlpackage ./CoreMLRunner/Sources/CoreMLRunner/

xcrun coremlcompiler compile コマンドは、.mlpackage ファイルを .mlmodelc ファイルにコンパイルするためのコマンドです。

プログラム側では .mlmodelc の URL を MLModel のコンストラクタに渡すだけです。

let configuration = MLModelConfiguration()
let modelURL = Bundle.module.url(forResource: "gpt2-baseline", withExtension: "mlmodelc")!
let model = try MLModel(contentsOf: modelURL, configuration: configuration)

また、MLModel の modelDescription から入力の shape などを取得可能です（該当箇所のソースコード）。

トークナイザの準備

swift-transformers では、Python の transformers と同様に、トークナイザが用意されています。

let tokenizer = try await AutoTokenizer.from(pretrained: "gpt2")

トークナイザの実装は複雑で、様々なアルゴリズムや正規化、前処理の手順が含まれています。これらを Swift で独自に実装するのは大変な作業ですが、swift-transformers が提供してくれているおかげで、Swift アプリでも、Python の transformers と同じ感覚で使うことができます。

生成プロセスの実装

テキストの生成プロセスは swift-transformers の Generation プロトコルで貪欲法 (greedy decoding) と top-k サンプリングの実装が提供されています。

このプロトコルを実装するとgenerate(config:tokens:model:callback) メソッドが使えるようになるので、残りは model 引数に渡す「入力の MLTensor を受け取り、次のトークンの出現確率を表す MLTensor を返す関数」を用意するだけです（以下のコードでは predictNextTokenScores が相当）。

class Generator: Generation {
    ...
    @discardableResult
    func generate(config: GenerationConfig, tokens: InputTokens) async -> GenerationOutput {
        let output = await generate(
            config: config,
            tokens: tokens,
            model: predictNextTokenScores
        ) { tokens in
            self.streamer?.put(tokens)
        }

        streamer?.end()
        return output
    }
}

predictNextTokenScores 関数は、入力の MLTensor を受け取り、次のトークンの出現確率を表す MLTensor を返す関数です。この関数では、以下の処理を行います。

パディング処理: 入力トークン列をモデルの入力形状に合わせるために、パディング処理を行います。
attention mask の生成: attention mask は、パディングされた部分が無視されるようにするために使用されます。
prediction(from:) メソッド: Core ML モデルの prediction(from:) メソッドを使用して、推論を実行します。

let inputDictionary = [
    "inputIds": inputIds,
    "attentionMask": attentionMask,
]
let outputs = try! await model.prediction(from: inputDictionary)

詳しくは実際のソースコードを見ていただければと思います。

ストリーミングの実装

Python 版のプログラム同様、TTFT/TPS を計測したいので、transformers の TextStreamer を移植したものを Generator で利用できるようにしています（ソースコード）。

let streamer = PerformanceMetricsStreamer(tokenizer: tokenizer)
let generator = Generator(model: model, maxContextLength: maxContextLength)
generator.streamer = streamer

最後に Generator.generate(config:tokens:) を呼び出せば、生成されたテキストが逐次出力されます。

_ = await generator.generate(
    config: generationConfig,
    tokens: inputIds
)

実行

swift コマンドでプログラムを実行してみます。

swift run --package-path ./CoreMLRunner -- coreml-runner --max-length=106 "What is generative AI?"
...
What is generative AI?

We've known that generative AI is much more complex because generative AI is more concerned with information that we do not have yet. In many fields we have already considered that humans can add new information whenever we please, and we would use only those that we don>t expect.
 the idea of 'new information' is not a new one. Most computer scientists think that the world is more complicated if we are concerned with information that we do not currently be able to know. While
[Prompt]  => 6 tokens, latency (TTFT): 0.21 ms
[Extend]  => 100 tokens, throughput: 8.53 tokens/s

見事に動きました！🎉 CPU/GPU/ANE の消費電力も見てみましょう。

想定通り、Neural Engine も動いてますね！🥳　また、Python 版と比べて GPU の稼働も上がり、逆に CPU は下がってることが分かります。

パフォーマンス

...でも、処理速度は遅いですね 🤔 ちゃんと、Python 版を計測したときと同じ方法で計測してみましょう。

poetry run python ./inference-metrics.py --warm 3 -n 5 -- swift run --package-path ./CoreMLRunner -- coreml-runner --max-length=106 "What is generative AI?"
...

Average Metrics:
[Prompt]  => 6 tokens, latency (TTFT): 0.15 ms
[Extend]  => 100 tokens, throughput: 9.03 tokens/s

Python 版は約 45 tokens/s だったので、実に 5 倍近く遅い結果となってしまいました。 ¹⁶

これは「Neural Engine が遅い」というわけではありません。主な理由として以下の二点が考えられます:

全シーケンス長での計算: 入力が最大コンテキストサイズまでゼロ埋めされているため、実際の処理対象トークン数が少なくても、全シーケンス長に対してアテンション計算が行われます。これにより、不要な計算が増加し、処理速度が低下します。

キー・バリューの再計算: 生成される各トークンに対して、以前に計算されたキー・バリュー対が再度計算されます。通常、これらの値はキャッシュ（KVキャッシュ）され、新しいトークンの処理時に再利用されますが、ベースラインモデルではこのキャッシュが使用されていないため、同じ計算が繰り返され、効率が悪化しています。

試しに、Python 版でも use_cache=False にして計測を再実行してみましょう。

$ poetry run python ./inference-metrics.py --warm 3 -n 5 -- poetry run python ./gpt-2/baseline/python-generate.py --max-length=106 --no-use-cache "What is generative AI?"
...
[Prompt]  => 6 tokens, latency (TTFT): 0.23 ms
[Extend]  => 100 tokens, throughput: 17.37 tokens/s

キャッシュを無効にするだけでパフォーマンスが格段に低下しています。

CoreML では、Stateful Models に対応することで、transformers のカスタム Cache を CoreML に載せることができます。

Apple Machine Learning Research のブログ記事 On Device Llama 3.1 with Core ML に、Llama 3.1 で KV キャッシュに対応する例が紹介されています。この記事では、KV キャッシュ以外にも可変長の入力による計算量削減、Int4 量子化を取り入れることで、ベースラインと比較して 177 倍の高速化 (0.19 tokens/s → 33.67 tokens/s) を達成しています。¹⁷

We Are Hiring!

ABEJA は、テクノロジーの社会実装に取り組んでいます。技術はもちろん、技術をどのようにして社会やビジネスに組み込んでいくかを考えるのが好きな方は、下記採用ページからエントリーください！（新卒の方やインターンシップのエントリーもお待ちしております！）

careers.abejainc.com

特に下記ポジションの募集を強化しています！ぜひ御覧ください！

Neural Engine は iPhone 8 / iPad（第 8 世代）/ M1 以降に搭載されています（ただし、CoreML から利用できるようになったのは 2018 年の iPhone XS/XR 以降）↩
macOS 15 以降、scaled_dot_product_attention（SDPA）が高レベルの操作としてサポートされています。この操作は単一の統合 GPU カーネルにマッピングされ、効率的に実行されます。これにより、非常に大きなサイズの「attention」テンソル（例: (1, #attn_heads, #token_length, #token_length)）を完全に生成する必要がなくなります（macOS 15 以前では、この操作は複数の処理に分解されていました）↩
モデルのステート管理のために、状態を入力・出力として明示する必要がなくなるためコードが簡潔化されるだけでなく、不要なデータのコピーが減り効率的な状態管理が可能です。特に Transformer モデルの KV (Key-Value) キャッシュを用いた推論では、余分な計算を回避し、高速化が実現されます↩
これまでも MLMultiArray や MLShapedArray のようなクラスはありましたが、これらはあくまで他のデータを変換するためのストレージであり、高度な操作は不可能でした↩
最も単純なのは、確率が最も高いトークンを次々に選ぶ貪欲法 (greedy decoding) です。しかし、この方法では文章が単調になりやすいため、他にもビームサーチや top-k サンプリング、top-p サンプリングといった戦略が用いられます。これらの戦略を選ぶ際には、生成される文章の多様性や文脈の整合性を考慮する必要があります。↩
Releasing Swift Transformers: Run On-Device LLMs in Apple Devices ↩
Git の preview ブランチから試すことができます。残念ながら、2024 年 12 月時点でも main にはマージされていません。↩
Settings > Access Tokens で Read access to contents of all public gated repos you can access を許可したトークンを発行します↩
https://x.com/karpathy/status/1777427944971083809 ↩
今回は、もっとも小さな 124M パラメータのモデルを使います。これくらい小さなパラメータ数であれば、強力な GPU がなくても気軽に試すことができます↩
できるだけ他のアプリケーションを起動していない状態で計測していますが、あくまで目安としての数値になります。↩
TorchScript は機能追加が凍結されており、将来的に非推奨となる予定です。しかし、代替手段と考えられる torch.export も実験段階であり、過渡期のあいだは両方を見ておく必要がありそうです。torch.export による変換は公式ドキュメントを参考にしてください。今回は、実験途中まで torch.export がサポートされていることに気づかず、試す時間がありませんでした 💦↩
他にも macOS15 などが定義されていますが、これらは iOS18 のエイリアスです↩
exporters の README にも It's normal for the conversion process to output many warning messages and other logging information. You can safely ignore these. という記述があります。↩
今回は Swift Package Manager を使ってコンソールで開発していくので、Xcode を起動する必要はありません↩
簡単のため、デバッグビルドで試していますが、リリースビルドで実行してもほぼ同じ結果です↩
今回は実際に試すだけの時間がありませんでした 🤮↩

ABEJA Tech Blog

中の人の興味のある情報を発信していきます

swift-transformers で LLM を動かしてみた

CoreML で機械学習モデルを動かす

`swift-transformers` を試す

Mistral 7B モデルを動かす

`swift-transformers` で推論を実装する

Python で動かしてみる

CoreML モデルに変換

Swift で動かす

パフォーマンス

We Are Hiring!

CoreML で機械学習モデルを動かす

swift-transformers を試す

Mistral 7B モデルを動かす

swift-transformers で推論を実装する

Python で動かしてみる

CoreML モデルに変換

Swift で動かす

パフォーマンス

We Are Hiring!

`swift-transformers` を試す

`swift-transformers` で推論を実装する