Google、MLハードウェアとフレームワークの統合を簡素化する「PJRT」を発表：MLインフラの断片化を解決する取り組みの一環

Googleは、オープンソースの機械学習（ML）コンパイラエコシステム「OpenXLA」プロジェクトの一部として、「PJRT」をオープンソースとして公開した。PJRTは、MLハードウェアやMLフレームワークに依存しないMLコンパイラインタフェースやMLランタイムインタフェースを提供する。

[＠IT]

　Googleは2023年5月8日（米国時間）、オープンソースの機械学習（ML）コンパイラエコシステム「OpenXLA」プロジェクトの一部として、「PJRT」をオープンソースとして公開した。PJRTは、MLハードウェアやMLフレームワークに依存しないMLコンパイラインタフェースやMLランタイムインタフェースを提供する。

　Googleは、MLインフラがさまざまなフレームワーク、コンパイラ、ランタイムにわたって断片化しており、この問題がMLモデルの開発スピードや、移植性、効率性、本番運用に悪影響を与えていると認識している。OpenXLAは、あらゆる主要なMLフレームワークのモデルをコンパイルして最適化し、幅広いハードウェア上で効率的なトレーニングやサービス提供を支援する目的で開発が進められている。

　Alibaba、Amazon Web Services（AWS）、AMD、Apple、Arm、Cerebras、Google、Graphcore、Hugging Face、Intel、Meta、NVIDIAなどがOpenXLAププロジェクトに参加している。

　PJRTは、OpenXLAプロジェクトの一部として、MLハードウェアとフレームワークの統合を簡素化するというビジョンを実現する。MLフレームワークの「TensorFlow」と「JAX」の主要なインタフェースであり、同じくMLフレームワークの「PyTorch」を完全にサポートしている。また、TPU、GPU、CPUでワークロードを実行するためのOpenXLAエコシステムとよく統合されている。さらに、Google社内のほとんどの本番ワークロードで既定のランタイム実行パスとなっている。

　PJRTは、ツールチェーンに依存しないアーキテクチャを採用しているため、あらゆるハードウェア、フレームワーク、コンパイラで利用でき、独自機能を追加するための拡張性も備えている。

MLハードウェアとフレームワークの統合の簡素化

　Googleは、PJRTが実現する「MLハードウェアとフレームワークの統合を簡素化する」というビジョンを次のように説明している。

　「MLワークロードの実行は現在、ハードウェア、フレームワーク、ワークフローの組み合わせマトリックスに対応しており、これは主に、緊密な垂直統合によって実現されている」

　「こうした垂直統合の例として、TPUとGPUの特定のカーネル、TensorFlowとPyTorchの特定のトレーニングおよびサービスツールチェーンのそれぞれの組み合わせがある。こうしたカスタムの1対1統合は、完全に有効なソリューションだが、ロックインを促進し、イノベーションを阻害し、メンテナンスコストが高くつく。この断片化されたソフトウェアスタックの問題は、時間とともに複雑化する。さまざまなコンピューティングハードウェアをサポートする必要があるからだ」

　「現在、さまざまなMLハードウェアが存在し、将来的にもハードウェアの多様性が増すと予想される。われわれの目標は、MLインフラ（フレームワーク、コンパイラ、ランタイム）へのハードウェアの統合を容易にすることで、MLワークロードをハードウェア間で移植可能にするという課題を解決することだ」

ワークロードのシームレスな実行を可能に

　このビジョンをPJRTで実現するワークフローは、次のようになる（図1）。

1. ハードウェア固有のコンパイラとランタイムプロバイダは、PJRT APIを実装し、コンパイラフックとランタイムフックを含むプラグインとしてパッケージ化し、フレームワークに登録する
2. フレームワークは、ワークロードを実行するハードウェアをターゲットとした動的ライブラリとして、1つまたは複数のPJRTプラグインを検出し、ロードする
3. フレームワークからターゲットハードウェアに対してワークロードを実行する

図1：特定のハードウェアをターゲットとしてPJRT APIが実装され、コンパイラおよびランタイムプラグインとしてパッケージ化され、フレームワークに呼び出される（提供：Google）

緊密なエコシステム

　PJRTは、OpenXLAプロジェクトの基礎的な柱として、OpenXLAプロジェクト内のプロジェクトとよく統合されている。これらのプロジェクトには、「StableHLO」とOpenXLAコンパイラ（「XLA」「IREE」）が含まれる。

　StableHLOは、さまざまなMLフレームワークとMLコンパイラの間の移植性レイヤーとして機能する、MLモデルにおける高レベル演算（HLO）のための演算セットだ。XLAは、一般的なMLフレームワークのモデルをGPU、CPU、MLアクセラレータなどのハードウェア向けに最適化するMLコンパイラ。IREEは、MLモデルを統合IR（中間表現）に変換するコンパイラおよびランタイムだ。

　PJRTは、TensorFlowとJAXの主要なインタフェースであり、PyTorch/XLAを通じてPyTorchを完全にサポートしている。フレームワークとハードウェアMLインフラの組み合わせに起因する断片化を解決するためのハードウェアインタフェースレイヤーを提供する（図2）。

図2：PJRTは、フレームワークとハードウェアMLインフラの組み合わせに起因する断片化を解決するためのハードウェアインタフェースレイヤーを提供する。OpenXLAとよく統合されている（提供：Google）