C# と ONNX Runtime Generative AI (DirectML) を使って Phi-3 をローカルで動かす

Build 2024 では Windows などローカルのリソースを使って Generative AI を動かすという話が非常に多かったように、Keynote でも度々取り上げられた Phi-3 についても AWQ で 4-bit 量子化された DirectML で利用可能な ONNX モデルが公開されています。

セッションでも話がありましたが、Microsoft としては DirectML を使っておけば GPU / NPU の両方に対応できるようにするようなので、今後はローカルでの AI 利用は DirectML が主導権を握る可能性がありそうです。

現状 Hugging Face で公開されている DirectML に対応した Phi-3 の ONNX モデルは以下の 4 種類です。Phi-3 mini と Phi-3 medium の両方が利用可能になっていますが、残念ながら現時点では DirectML に対応した Phi-3 vision は公開されていないようですが、近いうちにリリースされるようです。

• Phi-3 mini
  • microsoft/Phi-3-mini-4k-instruct-onnx · Hugging Face
  • microsoft/Phi-3-mini-128k-instruct-onnx · Hugging Face
• Phi-3 medium
  • microsoft/Phi-3-medium-4k-instruct-onnx-directml · Hugging Face
  • microsoft/Phi-3-medium-128k-instruct-onnx-directml · Hugging Face

Phi-3 の ONNX モデルの詳細と各 ONNX Runtime を使った場合の性能については ONNX Runtime のブログで紹介されていたので、こちらも一緒に載せておきます。

ようやく本題に入っていきますが、今回 Phi-3 の DirectML 対応の ONNX モデルがリリースされたことで、Python 環境や NVIDIA GPU を用意することなく Windows 上で Phi-3 が簡単に使えるようになりました。

既に公式ドキュメントには WinUI と ONNX Runtime Generative AI を使って、ローカルで Phi-3 を実行するサンプルが用意されていますが、WinUI 上で使いたいかと言われると微妙なのでシンプルにコンソールアプリから使ってみます。

ONNX Runtime Generative AI については公式ドキュメントが用意されているので、基本的な API についてはこちらを参照しておくと安心です。ONNX Runtime は C++ / C# / Python 向けに用意されていますが、今回は C# を使って Phi-3 を動かしてみます。

サンプルコードは GitHub で公開されているので、こちらをそのまま動かすのもありです。

実行に必要な Windows 向けの ONNX Runtime は CUDA / DirectML / CPU の 3 種類が用意されていますが、今回は AMD の GPU 上で実行したいので DirectML 版をインストールします。性能としては CUDA が圧倒的らしいので、可能な限り CUDA を使うべきなのでしょうが NVIDIA の GPU が無い場合は仕方ありません。DirectML を使っておくと NPU 対応PC では高速に動作するのが期待できます。

残念ながら現時点では Arm64 向けの ONNX Runtime Generative AI は開発中らしく、Windows Dev Kit 2023 上ではパフォーマンスが期待できません。ハードウェアアクセラレーションとして QNN に対応していないため、NPU を使うことも出来ず厳しそうです。新しい Surface が出るまでには対応されると信じたいですね。

今回作成したコンソールアプリでのサンプルは以下のようになります。基本的な流れは WinUI のサンプルから持ってきましたが、IAsyncEnumerable<T> や WinUI に依存する部分はバッサリ削除して ONNX Runtime Generative AI を使う部分に特化しています。Phi-3 のモデルについては Hugging Face からダウンロードしておきます。手順は書きませんが CLI や Git 経由でダウンロードできます。

using System.Diagnostics;

using Microsoft.ML.OnnxRuntimeGenAI;

var systemPrompt = "You are a helpful assistant.";
var userPrompt = "Microsoft について説明してください";

var prompt = $"<|system|>{systemPrompt}<|end|><|user|>{userPrompt}<|end|><|assistant|>";

Console.WriteLine("Loading model...");
var sw = Stopwatch.StartNew();
var model = new Model(@".\Phi-3-medium-128k-instruct-onnx-directml\directml-int4-awq-block-128");
var tokenizer = new Tokenizer(model);
sw.Stop();
Console.WriteLine($"Model loading took {sw.ElapsedMilliseconds} ms");

var sequences = tokenizer.Encode(prompt);
var generatorParams = new GeneratorParams(model);

generatorParams.SetSearchOption("max_length", 2048);
generatorParams.SetInputSequences(sequences);
generatorParams.TryGraphCaptureWithMaxBatchSize(1);

using var tokenizerStream = tokenizer.CreateStream();
using var generator = new Generator(model, generatorParams);

while (!generator.IsDone())
{
    generator.ComputeLogits();
    generator.GenerateNextToken();

    Console.Write(tokenizerStream.Decode(generator.GetSequence(0)[^1]));
}

基本的な流れとしては思ったよりシンプルで Model クラスで ONNX モデルを読み込み、次に Tokenizer の作成とプロンプトのエンコードを行い、最後に Generator を使ってトークンを生成するという流れです。

このコードを実行してみると、以下のように Phi-3 を使った回答が生成されました。この例では Phi-3 medium を利用して生成していますが、回答の生成はかなり早い部類と言えそうでした。

Phi-3 は主に英語向けに作られていますが、Phi-3 medium ではそれなりの日本語が生成されているので驚きです。ちなみに Phi-3 mini は完全に破綻した日本語が生成されるので、用途はそれなりに限定されそうです。

ちなみに推論中の GPU 負荷はそれなりに高く、モデルは可能な限り GPU メモリに読み込まれるようなので、以下のように推論中は GPU メモリがかなり必要になります。

Build のセッションでは GPU の場合メモリ帯域がかなり重要となると話がありましたが、この実行結果を見て納得出来ました。ちなみに 32GB メモリを搭載した Surface Laptop 4 でも CPU 内蔵の Intel Iris Xe で Phi-3 mini ぐらいであれば遅いですが動作しました。このあたりで DirectML の効果を実感しますね。

手持ちの Radeon RX 7900 XTX で Phi-3 medium を利用した場合は、平均して大体 70 token/sec ぐらいの性能でした。RDNA 3 で追加された AI Accelerator が使われているのか判断が付かないのが残念ですが、まだまだ最適化の余地はあると思います。

ちなみに Phi-3 mini の場合は 160 token/sec ぐらいだったので、Phi-3 medium の倍以上速くなっています。

Surface Laptop 4 の CPU 内蔵の Intel Iris Xe の場合は Phi-3 medium で 2 token/sec ぐらいだったので、新しい NPU 搭載の Surface ではどのくらいの性能になるのか今から楽しみです。