XLA（Accelerated Linear Algebra）は、機械学習用のオープンソース コンパイラです。XLA コンパイラは、PyTorch、TensorFlow、JAX などの一般的なフレームワークからモデルを取得し、GPU、CPU、ML アクセラレータなどのさまざまなハードウェア プラットフォームで高パフォーマンスの実行を実現するようにモデルを最適化します。たとえば、BERT MLPerf 提出では、8 Volta V100 GPU で XLA を使用すると、XLA を使用しない同じ GPU と比較して、パフォーマンスが約 7 倍向上し、バッチサイズが約 5 倍向上しました。
OpenXLA プロジェクトの一部として、XLA は、Alibaba、Amazon Web Services、AMD、Apple、Arm、Google、Intel、Meta、NVIDIA など、業界をリードする ML ハードウェアおよびソフトウェア企業によって共同で構築されています。
主なメリット
どこにでも構築可能: XLA は TensorFlow、PyTorch、JAX などの主要な ML フレームワークにすでに統合されています。
どこでも実行: GPU、CPU、ML アクセラレータなどのさまざまなバックエンドをサポートし、より多くのサポートを追加するプラグイン可能なインフラストラクチャを備えています。
パフォーマンスの最大化とスケーリング: 本番環境でテストされた最適化パスとモデル並列処理の自動パーティショニングにより、モデルのパフォーマンスを最適化します。
複雑さの排除: MLIR を活用して優れた機能を単一のコンパイラ ツールチェーンに集約するため、さまざまなドメイン固有のコンパイラを管理する必要がなくなります。
将来に備える: XLA は、主要な ML ハードウェア ベンダーとソフトウェア ベンダーのコラボレーションによって構築されたオープンソース プロジェクトとして、ML 業界の最先端で運用できるように設計されています。
ドキュメント
XLA の詳細については、以下のガイドをご覧ください。XLA を使用するのが初めての場合は、XLA アーキテクチャから始めて、次にコードレビューをご覧になることをおすすめします。