異種エッジでの高速かつポータブルな Llama2 推論

 (secondstate.io)
2 ポイント 投稿者 GN⁺ 2023-11-14 | 1件のコメント | WhatsAppで共有
  • Python 中心の推論スタックにおける大きな依存関係とデプロイ負担を減らすため、Rust+Wasm Llama2 アプリを 2MB のバイナリとしてパッケージ化し、複数のデバイスで実行する
  • 実装は llama.cpp を Wasm 向けに適合させたもので、GGUF モデルファイル、WasmEdge の GGML プラグイン、WASI NN API によりローカルのハードウェアアクセラレーションを活用する
  • 例では Llama-2-7B-Chat GGUF の 5-bit 量子化モデルを使用し、低価格帯の M2 MacBook で約 25 tokens/s、Nvidia A10G で約 50 tokens/s の性能を示す
  • WasmEdge ランタイムはコンテナツールと併用できるため、エッジデバイス、オンプレミス、クラウドに同じ Wasm バイナリをデプロイしやすい
  • GGML ツールはすでに利用可能だが初期段階であり、より多くのハードウェア・OS プラグイン、llama.cpp 設定のサポート、複数の Wasm 互換言語での WASI NN API サポートが今後の課題

Rust+Wasm で Llama2 推論を実行

  • Rust+Wasm スタックは、AI 推論において Python の代替として利用できる
    • Rust+Wasm アプリは Python と比べてサイズが 1/100、速度が 100 倍になり得るほか、同一バイナリをコード変更なしで複数環境のハードウェアアクセラレーションに対応させられる
  • Second State は、Llama2 モデル推論をネイティブ速度で実行するシンプルな Rust プログラムを作成した
    • Wasm にコンパイルされたバイナリアプリは 2MB
    • この Wasm アプリは、異種ハードウェアアクセラレータを備えたデバイス間で移植可能
    • WasmEdge ランタイムは、クラウド環境で安全な実行環境を提供する
    • WasmEdge はコンテナツールと併用することで、複数のデバイス上でポータブルアプリをオーケストレーションし実行できる

llama.cpp と GGUF ベース

  • この取り組みは、Georgi Gerganov が作成した llama.cpp に基づいている
  • 元の C++ プログラムを Wasm 上で動作するよう調整している
  • モデルファイルには GGUF 形式を使用する

実行手順

  • Linux または Mac デバイスに WasmEdge と GGML プラグインをインストールする
curl -sSf https://raw.githubusercontent.com/WasmEdge/WasmEdge/… | bash -s -- --plugins wasmedge_rustls wasi_nn-ggml
  • 事前ビルド済みの Wasm アプリをダウンロードする
curl -LO https://github.com/LlamaEdge/LlamaEdge/…
  • 例では、Llama2 7B chat-tuned モデルを 5-bit weights で量子化した GGUF ファイルを使用する
curl -LO https://huggingface.co/second-state/Llama-2-7B-Chat-GGUF/…
  • WasmEdge で Wasm 推論アプリを実行し、GGUF モデルを渡すと、対話形式で質問を入力できる
wasmedge --dir .:. --nn-preload default:GGML:AUTO:Llama-2-7b-chat-hf-Q5_K_M.gguf llama-chat.wasm

モデル設定と性能例

  • llama-chat.wasm はコマンドラインオプションでモデルとのやり取りの方式を設定する
    • --ctx-size: プロンプトのコンテキストサイズ、デフォルト 512
    • --n-predict: 予測するトークン数、デフォルト 1024
    • --n-gpu-layers: GPU で実行するレイヤー数、デフォルト 100
    • --batch-size: プロンプト処理のバッチサイズ、デフォルト 512
    • --temp: サンプリング温度、デフォルト 0.8
    • --repeat-penalty: 反復トークンのペナルティ、デフォルト 1.1
    • --prompt-template: llama-2-chat、codellama-instruct、mistral-instruct-v0.1、chatml、deepseek-chat、deepseek-coder などをサポート
    • --log-stat: 統計を出力
  • コンテキスト長 2048、応答最大 512 トークン、統計出力を指定した実行例は次のとおり
wasmedge --dir .:. --nn-preload default:GGML:AUTO:Llama-2-7b-chat-hf-Q5_K_M.gguf llama-chat.wasm -c 2048 -n 512 --log-stat
  • 低価格帯の M2 MacBook では、LLM の応答はデフォルトでストリーミング出力され、約 25 tokens/s で生成される
    • ログ例では eval time 基準で 82 回実行時に 25.64 tokens/s を記録した
    • kv self size は 1024.00MB、compute buffer total size は 630.14MB と表示される
  • Nvidia A10G 搭載マシンでは約 50 tokens/s で動作する
    • CUDA GPU アクセラレーションを使用し、35/35 レイヤーを GPU にオフロードする
    • VRAM 使用量はモデル 4474.93MB、コンテキスト 1648.02MB、合計 6122.95MB と表示される
    • eval time 基準で 83 回実行時に 50.55 tokens/s を記録した

OpenAI 互換 API サーバーと LLM アプリ

  • Rust と WasmEdge で作られた OpenAI 互換 API サーバーも提供されている
  • このサーバーを使うと、flows.network のような OpenAI 互換の開発者ツールで LLM エージェントやアプリを作成できる
  • 関連する説明は Build a super lightweight AI agent にある

Python 推論スタックの限界

  • Llama2 のような LLM は通常、Python ベースのフレームワークである PyTorch、Tensorflow、JAX で学習される
  • しかし AI コンピューティングの約 95% を占める推論アプリケーションでは、Python は不利だと見られている
  • Python パッケージは複雑な依存関係により、設定や利用が難しくなることがある
  • Python または PyTorch の Docker イメージは通常、数 GB から数十 GB にまで大きくなり得るため、エッジサーバーやデバイスでの AI 推論では負担が大きい
  • Python は C、C++、Rust のようなコンパイル言語より非常に遅く、最大で 35,000 倍遅くなる可能性がある
  • 実際のワークロードの多くは、Python ラッパー配下のネイティブ共有ライブラリに委譲される
    • この構造はデモには適しているが、ビジネス固有の要件に合わせて内部を修正するのは難しいと見られている
  • ネイティブライブラリへの高い依存と複雑な依存関係管理は、デバイス固有のハードウェア機能を活用する Python AI プログラムの移植性を低下させる
  • LLM ツールチェーンでよく使われる Python パッケージ同士が、pydantic のバージョン要件で衝突する例もある
    • llama-cpp-python は pydantic 2.0.1 を要求し、<=2.0 では動作しない
    • fastapi と chromadb は pydantic 1.9.2 を要求し、>=2.0 では動作しない

Rust+Wasm の利点

  • Rust+Wasm スタックは、デバイス、エッジクラウド、オンプレミスサーバー、パブリッククラウドをまたぐ統合クラウドコンピューティングインフラとして利用できる
  • AI 推論アプリケーションでは、Python スタックの強力な代替になり得る

  • 超軽量

    • 推論アプリケーションはすべての依存関係を含めて 2MB
    • 一般的な PyTorch コンテナサイズの 1% 未満

  • 高速実行

    • 前処理、テンソル計算、後処理の全体でネイティブ C/Rust の速度を出せる

  • 移植性

    • 同じ Wasm バイトコードアプリケーションが主要なコンピューティングプラットフォームで動作する
    • 異種ハードウェアアクセラレーションにも対応する

  • 設定・開発・デプロイの容易さ

    • 複雑な依存関係が減る
    • ノート PC 上で標準ツールを使って単一の Wasm ファイルをビルドし、複数環境にデプロイできる

  • 安全性とクラウド対応

    • Wasm ランタイムは信頼できないユーザーコードを隔離するよう設計されている
    • コンテナツールで管理し、クラウドネイティブプラットフォームにデプロイできる

Rust 推論プログラムの構造

  • デモ推論プログラムは Rust で書かれ、Wasm にコンパイルされる
  • 主要な Rust ソースコードは約 40 行
  • Rust プログラムは次の役割を担う
    • ユーザー入力の管理
    • 会話履歴の追跡
    • テキストを Llama2 chat template に変換
    • WASI NN API を使って推論を実行
  • 自分でビルドするには、Rust コンパイラと wasm32-wasi コンパイルターゲットをインストールする
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
rustup target add wasm32-wasi
  • その後、ソースプロジェクトをダウンロードし、cargo で Wasm ファイルをビルドする
git clone https://github.com/second-state/llama-utils
cd llama-utils/chat/
cargo build --target wasm32-wasi --release
cp target/wasm32-wasi/release/llama-chat.wasm .

クラウドとエッジへのデプロイ

  • Wasm バイトコードファイルがあれば、WasmEdge ランタイムをサポートするあらゆるデバイスにデプロイできる
  • 現在の GGML プラグインは generic Linux と Ubuntu Linux をサポートしている
    • x86 および ARM CPU
    • Nvidia GPU
    • Apple M1/M2/M3
  • WasmEdge GGML プラグインは llama.cpp をベースに、デバイスのハードウェアアクセラレーションを自動的に活用する
    • Nvidia GPU がある場合、インストーラが CUDA 最適化済みの GGML プラグインを自動でインストールする
    • Mac OS 向け GGML プラグインは Metal API を使い、M1/M2/M3 の内蔵 neural processing engine で推論ワークロードを実行する
    • Linux CPU ビルドは OpenBLAS を使い、AVX や SIMD など最新 CPU の計算機能を自動検出して活用する
  • この方式は、性能を犠牲にせずに異種 AI ハードウェアとプラットフォーム間の移植性を確保しようとするアプローチである

今後の作業

  • WasmEdge GGML ツールは現在利用可能で、クラウドネイティブの顧客に使われているが、まだ初期段階にある
  • 貢献できる領域は次のとおり
    • より多くのハードウェアおよび OS プラットフォーム向けの GGML プラグイン追加
      • TPU、ARM NPU、Linux および Windows の特殊な AI チップに関心がある
    • より多くの llama.cpp 設定のサポート
      • 現在は一部の設定オプションのみを Wasm から GGML プラグインへ渡せる
      • 目標は GGML が提供するすべてのオプションをサポートすること
    • 他の Wasm 互換言語での WASI NN API サポート
      • Go、Zig、Kotlin、JavaScript、C、C++ に関心がある

LLM 以外のモデルサポート

  • WasmEdge と WASI NN は、軽量・高速・移植可能・セキュアな Python 代替として、LLM 以外の人気 AI モデルに基づく推論アプリも作成できる
  • mediapipe-rs は、Google mediapipe Tensorflow モデルファミリー向けの Rust+Wasm API を提供する
  • WasmEdge YOLO は、YOLOv8 PyTorch モデルを扱う Rust+Wasm API プロジェクトである
  • WasmEdge ADAS demo は、Intel OpenVINO モデルで自動運転車の道路セグメンテーションを行う例である
  • WasmEdge Document AI は、OCR と文書処理モデルファミリー向けの Rust+Wasm API を提供する予定である
  • WasmEdge に関する議論や貢献は WasmEdge Discord で行える

関連記事

1件のコメント

 
GN⁺ 2023-11-14
Hacker News のコメント

  • Rust と WASM は良いが、コードを見ると基本的な Rust のコマンドラインスクリプト 150行にすぎない
    重い処理はモデルを WASI-NN バックエンドに渡す1行が担っており、ここでは WasmEdge ランタイムが提供している
    ただし WasmEdge は Rust ではなく C++で、この場合 Rust がもたらす利点はほとんどなく、バックエンドは Python など他の言語からも呼び出せる

    • ここで Rust がもたらす利点はバンドルと配布に見える
      Python と PyTorch をエンドユーザーがダブルクリックして実行できる形にまとめるのは今でもひどく面倒で、実際の高性能コードはどちらも C++だが、その C++にいくつか命令を送るために 2GB 以上と数万個のファイルを配布する代わりに、2MB の実行ファイル1つで済むなら意味がある

  • すごい仕事だ。ただ、試したい人は依然としてサイズの大きい重みファイルをダウンロードしなければならない
    本質的には、完全に移植可能で依存関係のない llama.cpp を 2MB にしたようなものだ
    アプリ開発者なら、配布可能なファイルの中に推論エンジンをパッケージングする最も簡単な方法かもしれない。重みはすでに移植可能で、必要なときにダウンロードすればよいので、本当に固定したい部分は推論エンジンだ

    • タイトルには 2MB の wasmと書いたほうが助けになりそう
      言われている通り、重みがそのサイズを圧倒する
    • 私のマシンで llama.cpp がビルドする main ファイルは 1.2MBだ
      2MB というサイズ自体は特に印象的ではなく、wasm をターゲットにすることでより移植可能になる点が肝で、特別に圧縮されているわけではない

  • これは結局 llama.cpp をラップしているだけでは？ 正直、x.cpp をラップするプロジェクトにはかなりうんざりしている
    この6か月間、Rust + WebGPU の機械学習フレームワークを開発してきて、GG の仕事がどれほどすごいかすぐに分かった
    まだ初期段階だが、ここで見られる:
    https://www.ratchet.sh/
    https://github.com/FL33TW00D/whisper-turbo

    • 何が印象的だったのか、もう少し説明してもらえる？ この分野をまったく知らないので、価値をきちんと理解するのが難しい
    • 他人の仕事をけなしながら、同じコメントで自分の仕事を宣伝したということ？ 倫理的に真剣に考え直す必要がある
    • GG って誰？

  • これが依存している wasm-nn、つまり https://github.com/WebAssembly/wasi-nn は、任意のチャンクをベンダー実装に送る方式の提案だ。API も実質的には入力設定、計算、出力設定という程度だ
    だからまったく移植可能ではない
    これが動作する理由は、llama.cpp にすでに実装されている抽象化に依存しているからだ。gguf モデルを複数のハードウェアターゲットへマッピングするそのコードを WasmEdge がそのまま取り込んでいる様子がここで見える: https://github.com/WasmEdge/WasmEdge/tree/master/plugins/was...
    したがって「開発者はバインディングを使って高水準言語で機械学習アプリを書き、WebAssembly にコンパイルしてから WasmEdge のような wasi-nn 対応ランタイムで実行できる」といった説明は完全に間違っている。実際にはそうはできない
    これは移植可能でもなく、サンドボックスでもなく、ハードウェア抽象化レイヤーでもない
    wasm バイナリがあっても、使っているランタイムのバージョンが偶然にも必要な特定の ggml バックエンドを実装している場合にだけ実行可能で、そうする要件はまったくないので、たいていはそうではない可能性が高い
    できたとしても、結局は llama.cpp の ggml コードを呼び出すことになるので、安全性はそのライブラリ次第だ
    記事では「移植性」と「Rust」を強調しすぎているが、実際にはどちらの利点もあまり見えてこない
    新しいハードウェアに WASI ランタイムがあると仮定して、モデルを実行できるのか？ GPU サポートはあるのか？ 答えは「そのプラットフォームで llama.cpp が GPU サポート付きでコンパイルされるか、使っているランタイムに ggml プラグインがあるか、その中に適合するバージョンの ggml がベンダリングされているかを確認してみて。そうでなければ無理」ということになる
    それなら、いったいなぜ WASI を使うのか？
    クロスプラットフォームの GPU サポートが難しいのは確かだが、これはかなりばかげて見える
    WebGPU が「GPU にバイナリチャンクを投げると、それが現在のハードウェアに合ったチャンクだった場合に、もしかしたら何か描画してくれるかもしれない」というような方式だったと想像してみれば、まさにそういう構造だ

    • 説明ありがとう。WASM でGPU サポートをどこから得ているのか気になっていた
    • セキュリティ面への影響も、もう少し詳しく説明してもらえる？

  • これを iPhone でオフライン実行できる？ そうなら、受信状況に関係なく基本的なインターネット検索のように使えそうで、キャンプのときに便利そうだ

    • さまざまな Linux、Mac、Windows ベースのデバイスで実行でき、Raspberry Pi や多くのノートPC・サーバーも含まれる
      ただしモデル自体を載せるには、依然として数GBのメモリが必要だ
    • Pixel でこのプロジェクトを動かしてみた。一部の iPhone/iPad でも動作するようだ
      [0] https://github.com/mlc-ai/mlc-llm
    • その用途なら Wikipedia の版をダウンロードするほうがよい可能性が高い。エントロピーのような問題もあるし
    • その用途である程度成功した商用 iOS アプリを作った
      最初は ggml で作ったが、その後 mlc-llm を見つけて、それをベースに移植した
      [0]: https://apps.apple.com/us/app/private-llm/id6448106860

  • 流れとして、エッジでTransformerアーキテクチャをより効率的かつ高速に実行する方法は今後も出てくるだろうが、VRAM要件はRustでは解決できないので、限界に近づいているように思う
    十分に大きなモデルを載せるときの主なボトルネックはまさにこの部分
    「小さなモデルは良くなっている。MistralやLlama 2を見よ」とは言えるが、小さなモデルも容量の限界に近づいている。70億パラメータに詰め込める情報には限界がある
    こうしたAIのアプローチがAGIにつながるとは思えない。あまりに非効率に感じる

    • 小さなモデルでも、MoEシステムやLoRAを動的に読み込む方式など、まだかなり余地があると思う

  • 「Mac OS向けGGMLプラグインはMetal APIを使い、M1/M2/M3内蔵のニューラル処理エンジンで推論処理を実行する」という説明は正確ではないように思う
    GGMLがMetal APIを使うというのは、M1/2/3のGPU上で実行されるという意味であって、Neural Engine上で実行されるという意味ではない
    もちろんそれ自体でも良いことだが、厳密にはそういうこと

  • コアライブラリを包んだラッパーが別の言語を使っているという理由で、プロジェクトが他の解法よりサイズを1/100にしたり、速度を100〜35000倍にしたりしたかのように語るクリックベイト的マーケティングが嫌い
    他の解法の周辺に積み上がったツールやコミュニティの専門性も完全に無視している
    まず、このプロジェクトはllama.cpp[1]ベースであり、数GB規模のモデルファイルをGPU/CPUに載せて実行する重い処理はllama.cppが担っている
    推論速度はラッパーの選択で制限されるものではなく、Go、Python、Node、Rustのラッパーもあるし、llama.cppを直接使ってもよい
    バイナリサイズもそれほど重要ではない。一般的な量子化モデルファイルは5GB〜40GBの範囲で、強力なGPUや16〜64GB RAMを備えたマシンが必要になるから
    [1] https://github.com/ggerganov/llama.cpp

  • サイズの大部分が結局、学習済みモデルの重みだとしたら、精度を失わずにどうやってサイズを数桁も削減できるのか？

    • 精度を失わずにサイズを減らすのは難しいと思う。ただし、量子化されたGGUFは素晴らしい
      ここで言う2MBは、モデルを除いたプログラムサイズを指しているように見える。wasmとllama.cppを実行するRustサーバーでllama.cppを動かす方式のようだ
      小さな llama.cpp/examples/server が好きでFreeChatに埋め込んでいるが、ツールの選択肢が増えるのはいつでも良いこと
      確認してみると、私が埋め込んでいるarm64/x86実行ファイルは現在4.2MB。FreeChatは12.1MBだが、デフォルトモデルが約3GBなので、2MBの差はあまり気にしていない
      [0]: https://github.com/ggerganov/llama.cpp/tree/master/examples/...
    • モデルそのもの、つまり学習済み重みのサイズを精度損失なしに数桁減らすという話なら、それは別の難しい課題
      この記事は推論アプリのサイズを100分の1にする話を扱っている

  • これをllama.cppを直接使う代わりに使う理由が分からない

    • ヒント：Rewrite-it-in-Rust経済圏の通貨は、実際に何かを動かすことではない
    • llama.cppは通常、OSとアーキテクチャごとにWindows、macOS、Linuxなど向けに別々にコンパイルする必要があり、移植性が低い
      記事では、異種ハードウェアアクセラレータを搭載したデバイスでハードウェアアクセラレーションを活用するとも説明している
      これは、wasmにコンパイルされたプログラムが、複数のデバイス上でGPUや専用AIチップのような多様なハードウェア資源を効率的に使えるという意味
      C++で直接実装すると、同程度の性能を出すにはハードウェアの種類ごとに別々の最適化やバージョンが必要になる場合がある