OpenXLA は、共同開発されたオープン ソース ML コンパイラ エコシステムです。
最近、機械学習分野の開発に特化した大手企業が、機械学習システム向けモデルのコンパイルと最適化のためのツールを共同開発することを目的としたOpenXLAプロジェクトを発表した。
このプロジェクトは、TensorFlow、PyTorch、および JAX フレームワークで準備されたモデルのコンパイルを統合して、さまざまな GPU および特殊なアクセラレータで効率的なトレーニングと実行を可能にするツールの開発を担当しています。 Google、NVIDIA、AMD、Intel、Meta、Apple、Arm、Alibaba、Amazon などの企業がプロジェクトの共同作業に参加しました。
OpenXLA プロジェクトは、複雑な ML インフラストラクチャの中で拡張できる最先端の ML コンパイラを提供します。 その基本的な柱は、ユーザーにとってのパフォーマンス、スケーラビリティ、移植性、柔軟性、および拡張性です。 OpenXLA を使用して、AI の開発と提供を加速することで、AI の真の可能性を解き放つことを目指しています。
OpenXLA を使用すると、開発者は主要なすべての ML フレームワークからモデルをコンパイルして最適化し、さまざまなハードウェアで効率的なトレーニングとサービスを提供できます。 OpenXLA を使用する開発者は、トレーニング時間、パフォーマンス、サービス レイテンシ、そして最終的には市場投入までの時間とコンピューティング コストが大幅に改善されます。
主要な研究チームとコミュニティ代表者の努力を結集することで、機械学習システムの開発を促進し、様々なフレームワークやチームにおけるインフラストラクチャの断片化の問題を解決できることが期待される。
OpenXLAは、機械学習モデルの基盤となるフレームワークに関わらず、様々なハードウェアに対する効果的なサポートの実装を可能にします。OpenXLAによって、モデルのトレーニング時間の短縮、パフォーマンスの向上、レイテンシの低減、計算オーバーヘッドの削減、そして市場投入までの時間の短縮が期待されます。
OpenXLAは主に3つのコンポーネントで構成されており、そのコードはApache 2.0ライセンスの下で配布されています。
- XLA (加速線形代数) は、さまざまなメーカーの GPU、CPU、特殊なアクセラレータなど、さまざまなハードウェア プラットフォームで高性能に実行できるように機械学習モデルを最適化できるコンパイラです。
- StableHLO は、機械学習システム モデルで使用する一連の高レベル操作 (HLO) の基本仕様および実装です。 これは、モデルを特定のハードウェアで実行できるように変換する、機械学習フレームワークとコンパイラの間のレイヤーとして機能します。 PyTorch、TensorFlow、および JAX フレームワーク用の StableHLO 形式でモデルを生成するためのレイヤーが用意されています。 MHLO スイートは StableHLO の基盤として使用され、シリアライゼーションとバージョン管理のサポートによって拡張されています。
- IREE (Intermediate Representation Execution Environment) は、LLVM プロジェクトの MLIR (Intermediate Multi-Level Representation) 形式に基づいて、機械学習モデルを汎用中間表現に変換するコンパイラおよびランタイムです。 機能のうち、プリコンパイル (事前) の可能性、フロー制御のサポート、モデルで動的要素を使用する機能、さまざまな CPU と GPU の最適化、および低オーバーヘッドが強調されています。
OpenXLAの主な利点としては、デバイス固有のコードを記述することなく最適なパフォーマンスが実現できることに加え、代数式の簡略化、効率的なメモリ割り当て、実行スケジューリングなど、すぐに使用できる最適化機能が提供され、最大メモリ消費量とオーバーヘッドの削減が考慮されている点が挙げられます。
もう一つの利点は、計算のスケーリングと並列化が簡素化されることです。開発者は、重要なテンソルのサブセットに注釈を追加するだけで、コンパイラが並列計算用のコードを自動的に生成できます。
また、 AMDやNVIDIAのGPU、x86やARMのCPU、Google TPU MLアクセラレータ、AWS Trainium Inferentia IPU、Graphcore、Cerebras Wafer-Scale Engineなど、複数のハードウェアプラットフォームをサポートすることで移植性も確保されている点も強調されています。
CUDA、HIP、SYCL、Tritonなどの並列コンピューティング言語を使用して深層機械学習のプリミティブを記述するサポートや、モデル内のボトルネックを手動で調整する機能など、追加機能の実装を伴う拡張機能の接続をサポートします。
最後に、さらに詳しく知りたい場合は、以下のリンクから詳細をご確認いただけます。