- Rust+WASMスタックは、AI推論においてPythonの強力な代替になり得る
- Pythonと比べてRust+Wasmアプリはサイズが1/100、速度が100倍速く、バイナリコードを変更しなくても完全なハードウェアアクセラレーションであらゆる場所で安全に実行できる
- RustはAGIの言語である
- Llama2推論をネイティブ速度で実行するシンプルなRustプログラムを書いてみると
- Wasmにコンパイル後のバイナリはわずか2MBだが、異種ハードウェアアクセラレータへ完全に移植可能
- また、Wasmランタイム（WasmEdge）はクラウド環境向けの安全なセキュア実行環境を提供する。コンテナツールとシームレスに連携し、さまざまなデバイスでポータブルなアプリケーションをオーケストレーションして実行できる
なぜPythonではないのか？
- llama2のようなLLMは一般にPython（例: PyTorch、Tensorflow、JAX）で学習される
- しかし、AIコンピューティングの約95%を占める推論アプリケーションにPythonを使うのは大きな誤り
- Pythonパッケージは複雑な依存関係を持っている。セットアップも利用も難しい
- Pythonの依存関係は非常に重い。PythonやPyTorch向けDockerイメージは通常数GBから数十GBに達し、これはエッジサーバーやデバイスでAI推論を行う際に特に問題になる
- Pythonは非常に遅い言語で、C、C++、Rustのようなコンパイル言語より最大35,000倍遅い
- Pythonは遅いため、実際のワークロードの大半はPythonラッパーの下にあるネイティブ共有ライブラリへ委譲しなければならない
- そのためPython推論アプリはデモには優れているが、業務ごとの要件に合わせて内部を修正するのは非常に難しい
- ネイティブライブラリへの依存度が高く依存関係管理も複雑なため、デバイス固有のハードウェア機能を活用しながらPythonのAIプログラムを複数のデバイスへ移植するのは非常に難しい
- LLMツールチェーンで一般的に使われるPythonパッケージは、互いに直接衝突する
- LLVM、TensorFlow、Swift言語で知られるChris Lattnerが今週スタートアップのポッドキャストで素晴らしいインタビューを行い、Pythonがモデル学習には優れていても推論アプリケーションには不適切な選択である理由を説明している
Rust+Wasmの利点
- Rust+Wasmスタックは、デバイスからエッジクラウド、オンプレミスサーバー、パブリッククラウドまでを一貫してつなぐクラウドコンピューティング基盤を提供する
- AI推論アプリケーション向けPythonスタックの強力な代替であり、イーロン・マスクがRustをAGIの言語と呼んだのも不思議ではない
- 超軽量: 推論アプリケーションはすべての依存関係を含めてもわずか2MB。一般的なPyTorchコンテナサイズの1%にも満たない
- 非常に高速: 前処理、テンソル演算、後処理など、推論アプリケーションのあらゆる部分でネイティブC/Rust並みの速度を提供する
- ポータブル: 異種ハードウェアアクセラレーションをサポートするすべての主要コンピューティングプラットフォームで、同一のWasmバイトコードアプリケーションを実行できる
- セットアップ、開発、デプロイが容易: もはや複雑な依存関係はない。ノートPCの標準ツールで単一のWasmファイルをビルドし、どこへでもデプロイできる
- 安全でクラウド対応: Wasmランタイムは、信頼できないユーザーコードを隔離するよう設計されている。Wasmランタイムはコンテナツールで管理でき、クラウドネイティブなプラットフォームへ容易にデプロイできる
Rust推論デモ
- 40行のRustコードで書かれ、Wasmにコンパイルされる
- Wasmバイトコードができれば、WasmEdgeランタイムをサポートするあらゆるデバイスへデプロイ可能
- llama.cppベースのWasmEdge GGMLプラグインは、デバイスのハードウェアアクセラレーションを自動的に活用してllama2モデルを実行する
次のステップ
- WasmEdge GGMLツール群はすでに利用可能で、実際にクラウドネイティブの顧客が使っているが、まだ初期段階にある
- オープンソースプロジェクトに貢献し、将来のLLM推論インフラの方向性づくりに関心があるなら、貢献できる可能性がある項目がいくつかある
- より多くのハードウェアおよびOSプラットフォーム向けGGMLプラグインを追加する: LinuxとWindowsのTPU、ARM NPU、そのほかの特殊AIチップなど
- より多くのllama.cpp構成をサポートする: 現在は一部の構成オプションをWasmからGGMLプラグインへ渡すことをサポートしているが、GGMLが提供するすべてのオプションに対応したい
- 他のWasm互換言語向けWASI NN APIをサポートする: 特にGo、Zig、Kotlin、JavaScript、C、C++
5件のコメント
現在AIで劣勢にあるAppleが取り得る最も現実的なアプローチは、
本文のような最適化されたS/W推論エンジンをiPhoneのAチップのNPUで高速化し、foundationモデルのファイルをiPhoneに丸ごと入れてしまうことのように思えます。
最近は学界・業界で、AI性能はサイズよりもデータおよびfine-tuningが重要だというのが大勢の見方であり、Llama2はfine-tuning性能がかなり良いという評価も多いため、十分に学習させたモデルを継続的にedge deviceへ配信することも十分可能だと思われます。そういう意味では、iPhoneにすべて詰め込む戦略はかなり現実的でしょう。
WASMはこんなに盛り上がっているんですか？ でも、PythonはMojoと比較すべきではないでしょうか
同意します。Mojoはまだ実験段階だと見るべきでしょう
Mojoはまだ正式公開されていませんが、Rustはすでにかなり成熟した言語であり、実務の現場でさまざまな部分を置き換えています。
GN⁺: 異種エッジでの高速かつポータブルな Llama2 推論
GN⁺ の要約と Hacker News のコメントもあわせて参考にしてください。