XLA

XLA(Accelerated Linear Algebra)は、ML 用のオープンソース コンパイラです。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 アーキテクチャから始めて、次にコードレビューをご覧ください。