LLaMA 2 エンドツーエンド推論が利用可能になりました。中国チームより

Buddy CompilerのエンドツーエンドLLaMA2-7B推論例がbuddy-mlirリポジトリ[1]のメインラインにマージされました。

Buddy Compiler のフロントエンドに TorchDynamo 用のサードパーティ コンパイラを実装し、 MLIR と PyTorch のコンパイル エコシステムを組み合わせました。

現在、フロントエンドは LLaMA 計算グラフをカバーできます。MLIR に変換した後、いくつかのベクトル化と並列最適化を統合し、AVX512 プラットフォームでテストしました。推論プロセス全体は実行できますが、依然として多くの最適化が必要です。関連するリンクと現在のステータスは次のとおりです。

• [E2E] バディコンパイラエンドツーエンドLLaMA2-7B推論例 [2]
• [E2E] 上記のエンドツーエンドの推論例の目的は、コンパイラスタックの設計を示すことであり、完全なLLaMAの質問回答ツールではありません。
• [フロントエンド] Buddy Dynamoコンパイラ[3]
• [Midend] 行列乗算のベクトル化とループの並列最適化を統合
• [バックエンド] X86 AVX512 マシン (Ubuntu 22.04) でテストされたエンドツーエンドの例
• [WIP] さまざまな最適化を開発して統合する (現時点では遅すぎる)
• [WIP] さまざまな CPU バックエンド (Arm Neon、RVV など) でのテスト
• [WIP] さまざまなハードウェア プラットフォーム (GPU、Gemmini Accelerator など) に接続します。
• [WIP] フロントエンドのカバレッジを拡張 (Stable Diffusion、CLIP、Whisper など)

概要

大規模 AI モデルの爆発的な増加により、ソフトウェアとハ​​ードウェアの設計に新たな手法と機会がもたらされました。モデルの規模、種類、モダリティが拡大し、多様化するにつれて、単一のソフトウェアおよびハードウェア テクノロジ スタックでさまざまなシナリオをカバーする能力はますます制限されるようになります。この傾向により、ソフトウェアとハ​​ードウェアのエコシステムの重要性についての理解も深まりました。私の意見では、エコシステム全体にとって最も重要な 3 つの設計原則は次のとおりです (最も重要なものから最も重要でないものの順)。

• 技術ルートの標準化（エコシステムの基盤）
• 参入障壁が低い（十分な貢献者とプレイヤーが生態系の繁栄の鍵となる）
• 高い最適化上限（エコシステムの開発可能性を決定）

上記の考慮に基づいて、ソフトウェアとハ​​ードウェアの共同設計エコシステムの実現に特化した Buddy Compiler を作成しました。私たちの目標は、ドメイン固有プログラミング言語 (DSL) からドメイン固有ハードウェア アーキテクチャ (DSA) へのコンパイル フローと共同設計を可能にすることです。

この記事では、Buddy Compiler を使用して LLaMA 2 のエンドツーエンドの推論を完了する方法を紹介します。同時に、Buddy Compiler の全体的な設計コンセプトと将来の計画も共有します。具体的な構築プロセスについては、こちらのドキュメント[4]を参照してください（大規模なモデルを構築するには、多くの忍耐と優れたパフォーマンスを持つマシンが必要です）。PyTorchとMLIRに基づくLLaMAのエンドツーエンドの推論パスを開設しましたが、完全なパフォーマンス分析と最適化はまだ実施していません。

現在、行列乗算にはベクトル化の最適化、ループには並列計算の最適化のみを適用しています。

最適化とチューニング戦略はまだ開発中であるため、パフォーマンスは現時点では低いレベルにあります。当社の当初の目標は、極端なパフォーマンスを追求することではなく、標準的なエンドツーエンドのチャネルを確立し、それに基づいてあらゆるレベルでパフォーマンスを最適化および調整することです。

テクニカルルート

技術ルートの標準化は、私たちが目指す中核原則であるだけでなく、これによりより多くの貢献者やパートナーを引き付けると同時に、ユーザーの学習のハードルとその後のメンテナンスコストを効果的に削減できると確信しています。

図に示すように、さまざまな AI モデルを接続するために PyTorch エコシステムを選択し、Buddy Compiler の中間表現として MLIR エコシステムを選択します。

私たちは、さまざまな AI モデルを複数のハードウェア アーキテクチャにマッピングするというビジョンを実現することを目指して、Buddy Compiler を Torch Dynamo のカスタム コンパイラとして使用してコンパイル スタック全体を形成します。設計上のポイントは次のとおりです。

AIモデルに接続するためのトレースツールとしてTorchDynamoを使用する

TorchDynamoはCPythonのフレーム評価API機能を使用して、PyTorchグラフをより正確にキャプチャします。詳細については、PyTorchのドキュメント[5]を参照してください。さらに、PyTorch 2.xは包括的なコンパイルサポート、特に非常にユーザーフレンドリーなカスタムコンパイラドッキングソリューションを提供します[6]。したがって、PyTorch 2.x TorchDynamo が AI モデルに接続するための唯一の選択肢になりました。

ドッキングレベルとしてAten IRを選択します

PyTorchのドキュメント[5]によると、Core Aten IRはバックエンドにサービスを提供する演算子のセットであり、Prime IRよりも高いレベルの抽象化を備えています。 MLIR へのマッピングには、より高度な抽象化レベルを使用することをお勧めします。これにより、多層最適化のためのスペースと情報が増えます。そのため、Linalg Dialect と TOSA Dialect のさまざまな操作に簡単にマッピングできる Aten IR をドッキング レベルとして選択します。

MLIR Python バインディングを使用して Dynamo コンパイラを実装し、TOSA/Linalg Ops を生成する

Buddy Compiler フロントエンドの Dynamo Compiler (または Dynamo Importer) は、PyTorch の Aten IR を MLIR に変換するために使用され、モデル パラメータの処理とパッケージ化も行います。

Dynamo コンパイラは、TorchDynamo から渡された FX グラフをトラバースし、MLIR Python バインディングを使用して各 Aten IR ノードを変換し、ターゲット MLIR コードを生成します。

Dynamo コンパイラはドッキング方言の優先順位設定をサポートしています。つまり、変換の最初の選択肢として Linalg 方言または TOSA 方言を選択できます。

最後に、Dynamo コンパイラは、変換された MLIR モジュールと重みパラメータを出力する役割を担います。

Buddy Compilerツールチェーンを使用した最適化と降下

現在、私たちのコンパイル パイプライン全体は Python スクリプトを使用して完全に完了しているわけではなく、フロントエンド、ミドルエンド、バックエンドを統合するために CMake を使用しています。

これにより、フロントエンド、ミドルエンド、バックエンドの開発プロセスがある程度簡素化され、分離されます。コンパイルの最適化と下降のすべてのパスは、Buddy-opt ツールに登録されます。このツールには、Buddy コンパイラーのすべてのアップストリーム パスと最適化パスが含まれるため、buddy-opt はアップストリーム mlir-opt のスーパーセットになります。当社では、再利用性を最大限に高めるために、MLIR Core Dialect を使用して一般的なハードウェアのコンパイル プロセスを実装しています。そのため、当社のツールは、mlir-translate、llc などのすべての LLVM/MLIR ツールと互換性があります。

現在、ループには並列コンピューティングの最適化を使用しています。この最適化では、関連する中間表現と最適化パスが完全に上流の OMP 方言からのものであり、直接再利用して優れた最適化効果をもたらすことができます。このことから、統合されたエコシステムの利点もわかります。さらに、コンパイルの最適化のために粗粒度演算用のベクトル化アルゴリズムを設計しており、Vector Dialect を使用することで、クロス SIMD/Vector プラットフォーム効果も実現できます。

特定のアクセラレータ (Gemmini Accelerator など) 用のコードを生成する場合は、buddy-translate と buddy-llc を使用してハードウェア固有の LLVM IR を生成し、最終的にアクセラレータのハードウェア命令を生成することもできます。

フォワード関数は、バディコンパイラテキストコンテナと組み合わせてエンドツーエンドの推論を完了するために使用されます。

コンパイルの最適化と下降後、モデルの Forward 関数が共有ライブラリとして構築されます。

私たちのシナリオの多くはシミュレーターや開発プラットフォームでテストする必要があるため、実行プロセスを Python エコシステムに引き渡すのではなく、AOT コンパイルを実行して実行可能ファイルを生成しました。

MLIR から構築された Forward 関数を使用してエンドツーエンドの推論を実現するために、テキスト入出力用のトークナイザーおよびデータ コンテナーとして、Text Container と MemRef Container の C++ バージョンを提供します。

最後に、C++ のメイン関数で、入力テキストと重みパラメータがデータ コンテナーにロードされ、推論のために Forward 関数が呼び出されます。出力トークンは、最終的なテキストを取得するための後処理のためにテキスト コンテナーに送信されます。

今後の仕事

現在、LLaMA から CPU SIMD/Vector プラットフォームへのパスを開設し、X86 AVX512 を使用して予備テストを実施しています。 Vector Dialect のクロスプラットフォーム性のため、Arm Neon と RISC-V Vector Extension もネイティブでサポートされており、広範囲にわたるテストを実施しています。

同時に、Buddy Compiler の Gemmini アクセラレータと大規模モデル推論パスの接続もサポートしようとしています。

さらに、GPU最適化パスも開発中です。モデル レベルでは、より大規模なマルチモーダル モデルをテストして接続し、Buddy Dynamo Compiler フロント エンドの範囲をさらに拡大したいと考えています。上記の作業を完了した後の次のマイルストーンは、マルチモーダル大規模モデルを複数のハードウェア プラットフォームにコンパイルすることです。

一般的に、フロントエンド、ミドルエンド、バックエンドは比較的分離された状態で開発され、最終的には接続されます。

• フロントエンド: より多くのモデルを接続して、オペレーターのカバレッジを向上させます。
• 中規模: より詳細なパフォーマンス分析を実施し、さまざまなコンピューティング負荷に対してさまざまなレベルの最適化を開発します。
• バックエンド: 各 CPU SIMD/ベクター プラットフォームをテストし、GPU/Gemmini アクセラレータのエンドツーエンド パスを完了します。

3 つの部分すべてが比較的完成したら、Python インターフェイスを使用してすべてのツール チェーンをパッケージ化し、よりエレガントな使用プロセスを形成することを検討してください。

要約する

現在、大規模モデル推論用のソフトウェア スタックが次々と登場しており、技術的なルートも異なります。私たちが使用する Torch 2.x + Buddy Compiler コンパイル スタック設計戦略は、実際には AI コンピューティング用の PyTorch と MLIR という 2 つの主要なエコシステムを統合することです。もちろん、MLIR エコシステムにおける大規模モデル推論の場合、Torch-MLIR + IREE は現時点では比較的完全なソリューションであると考えています。 nod.aiのSHARK Turbine[7]はこの技術的アプローチを採用しています。

Torch-MLIR と IREE の組み合わせと比較して、Buddy Compiler はシンプルでありながら強力な設計を重視し、極端な再利用戦略と完全なコード生成戦略を採用しています。 Torch-MLIR の Torch Dialect レベルと比較すると、Buddy Compiler は TOSA/Linalg を直接再利用して Aten IR に接続することを優先します。また、すべてのバックエンドをカバーする IREE のランタイムおよび実行エンジンの設計と比較すると、Buddy Compiler は完全なコード生成を実行することを優先します。

Shark Turbine [8] のカバーは航空機のタービンエンジンで、彼らの技術方針によく合っています。TorchDynamo + Torch-MLIR + IREE は、非常に洗練された高耐久性のコンパイルスタックです。このような「エンジン」があれば、理論的にはどんな山や渓谷も越えることができます。それに比べて、Buddy Compiler は、さまざまな個性の電気自動車を構築できる、電気自動車用の 3 電気プラットフォームのようなものです。私たちにとって、LLaMA のサポートは出発点でも終点でもありません。それは、探検の道で出会う山です。私たちは、その山を登り、山の向こう側にある世界を見て、自分たちで作った車を運転したいと思っています。

謝辞

Buddy Compiler のすべての貢献者、特に LLaMA の運営に尽力してくれたパートナーの zhanghb97、weilinquan、xTayEx、EllisLambda、Lester-1、LAJIidea、SForeKeeper、LHY-24、xlinsist、qingqing12138 に感謝します。また、オープンソース プロジェクトの席を提供してくれた OSPP 組織委員会にも感謝します。