Googleらが共同開発、AI／機械学習を支援するオープンソース技術「OpenXLA」とは：AIの開発シーンで生まれる悩みとは

AI/ML業界の多くの主要企業が共同開発しているオープンソースMLコンパイラエコシステム「OpenXLA」の利用と貢献が可能になった。

[＠IT]

　Googleは2023年3月8日（米国時間）、AI（人工知能）／ML（機械学習）業界の多くの主要企業が共同開発を進めているオープンソースMLコンパイラエコシステム「OpenXLA」（XLA：Accelerated Linear Algebra）の利用と貢献を可能にしたと発表した。

　OpenXLAの開発にはAlibaba、Amazon Web Services（AWS）、AMD、Apple、Arm、Cerebras、Google、Graphcore、Hugging Face、Intel、Meta、NVIDIAなどが参加している。

利用するフレームワーク増加に伴うサイロ化が課題に

　現在のML開発、導入では「TensorFlow」「PyTorch」「JAX」といったMLフレームワークや、多種多様なハードウェアの利用により、インフラが断片化、サイロ化していることが問題となっている。問題の背景には、現実世界における複雑な問題を解決するため、MLモデルのパラメーター数が指数関数的に増加していること、ディープラーニングの計算量が半年ごとに倍増していること、開発者がインフラのパフォーマンスの最大化と最大限の活用を追求していることがある。

　こうしたMLインフラの断片化、サイロ化は、開発スピードや、モデルの移植性、効率性、本番運用に悪影響を与える。OpenXLAは、あらゆる主要なMLフレームワークのモデルをコンパイルして最適化し、幅広いハードウェア上で効率的なトレーニングやサービス提供を支援する。Googleは、OpenXLAを活用することでトレーニング時間、スループット、サービスレイテンシ、さらには市場投入時間やコンピュートコストが大幅に改善されると述べている。

OpenXLAのアーキテクチャ

　MLインフラの断片化にはさまざまな要因が絡むことから、単一の組織だけでは、この問題を効果的に解決できない。Googleは「フレームワークからコンパイラ、ランタイム、半導体まで、AIスタックのさまざまなレベルで活動する開発者や主要企業が協力して取り組むOpenXLAプロジェクトは、断片化問題に対処するのに適している」と説明している。

　OpenXLAプロジェクトには、「XLA」「StableHLO」「IREE（Intermediate Representation Execution Environment：中間表現実行環境）」というモジュラーツールチェーンのリポジトリが含まれている。断片化問題に対処するため使われるこれらのモジュラーツールチェーンは、いずれもMLIR（マルチレベル中間表現）を利用する。MLIRは、MLモデルの一貫した表現、最適化、ハードウェア上での実行を可能にするコンパイラインフラだ。

OpenXLAの高レベルのコンパイルフローおよびアーキテクチャ。図中の最適化、フレームワーク、ハードウェアターゲットは、OpenXLAによって開発者に提供されるものの一部が示されている（提供：Google）

　これらのモジュラーツールチェーンの概要は以下の通り。

XLA

　GPU、CPU、MLアクセラレータ向けのオープンソースMLコンパイラ。XLAコンパイラは、PyTorch、TensorFlow、JAXなどの一般的なMLフレームワークのモデルを、GPU、CPU、MLアクセラレータなど、さまざまなハードウェアプラットフォームにわたって高パフォーマンスで実行されるように最適化する。

IREE

　MLIRベースのエンドツーエンドコンパイラおよびランタイム。MLモデルを統合IRに変換する。統合IRは、データセンターのニーズに合わせてスケールアップし、モバイルやエッジ環境での制約や特別な考慮事項に応じてスケールダウンする。

StableHLO

　MLモデルにおける高レベル演算（HLO）のための演算セット。基本的に、さまざまなMLフレームワークとMLコンパイラの間の移植性レイヤーとして機能する。StableHLOプログラムを生成するMLフレームワークは、StableHLOプログラムを利用するMLコンパイラと互換性がある。

　StableHLOは、さまざまなMLフレームワーク（TensorFlow、JAX、PyTorchなど）とMLコンパイラ（XLA、IREEなど）の間の相互運用性を高めることで、ML開発を簡素化、迅速化することを目指している。主要なMLフレームワークであるensorFlow、JAX、PyTorchは、全てStableHLOを生成できる。StableHLOチームは、2023年を通じてPyTorchチームと緊密に協力し、「PyTorch 2.0」への統合を可能にする計画だ。

　この3つのモジュラーツールチェーンを含むOpenXLAは、以下のようなメリットを提供する。

• 幅広いMLユースケースに対応する
• モデルのパフォーマンスを簡単に高速化できる
• 最小限の労力でワークロードをスケーリングできる
• さまざまなGPU、CPU、MLアクセラレータなど、多数のハードウェアデバイスをサポートし、移植性レイヤーであるStableHLOを介して、TensorFlow、PyTorch、JAXをサポートする
• さまざまなカーネル言語でディープラーニングプリミティブを記述でき、ハードウェア機能を最大限に活用できる