[論文レビュー] LazyTensor: combining eager execution with domain-specific compilers

Alex Suhan, Davide Libenzi+5 more|arXiv (Cornell University)|Feb 26, 2021

Parallel Computing and Optimization Techniques参考文献 28被引用数 6

ひとこと要約

LazyTensor は、define-by-run ML フレームワークにおいてホスト言語の表現力をフルに活用しつつ、ドメイン特化コンパイラ（例：XLA）によるテンソル演算のコンパイルを可能にすることで、言語サブセット問題を回避する。これは、演算を遅延実行可能で合成可能な中間表現（IR）グラフとして表現し、必要に応じてのみコンパイルする仕組みにより実現され、PyTorch および Swift for TensorFlow における TPU や GPU でも高いパフォーマンスを発揮する。言語やハードウェアを跨ぐ共有実装が可能である。

ABSTRACT

Domain-specific optimizing compilers have demonstrated significant performance and portability benefits, but require programs to be represented in their specialized IRs. Existing frontends to these compilers suffer from the "language subset problem" where some host language features are unsupported in the subset of the user's program that interacts with the domain-specific compiler. By contrast, define-by-run ML frameworks-colloquially called "eager" mode-are popular due to their ease of use and expressivity, where the full power of the host programming language can be used. LazyTensor is a technique to target domain specific compilers without sacrificing define-by-run ergonomics. Initially developed to support PyTorch on Cloud TPUs, the technique, along with a substantially shared implementation, has been used by Swift for TensorFlow across CPUs, GPUs, and TPUs, demonstrating the generality of the approach across (1) Tensor implementations, (2) hardware accelerators, and (3) programming languages.

研究の動機と目的

• ドメイン特化コンパイラにおける言語サブセット問題（ユーザーのプログラムが制限された IR 表現力に従わなければならないこと）を解決すること。
• define-by-run ML フレームワークが、デバッグ性や言語の柔軟性を損なわず、高性能かつハードウェア最適化されたコンパイラを活用できるようにすること。
• 複数のハードウェアアクセラレータ（CPU、GPU、TPU）およびテンソル型をサポートする、再利用可能なクロスランゲージ実装を構築すること。
• 遅延コンパイルが、HuggingFace Transformers などの実世界の ML ワークロード（Cloud TPU 上）で顕著なパフォーマンス向上を達成できることを示すこと。

提案手法

• テンソル演算を、即時実行中に構築されるが、必要に応じてのみコンパイルされる遅延実行可能で合成可能な IR グラフとして表現する。
• Python や Swift といった異なるプログラミング言語、および PyTorch や Swift for TensorFlow といった異なるフレームワーク間で、IR グラフの構築とコンパイルパイプラインを共有実装する。
• XLA などのドメイン特化コンパイラと統合し、必要に応じて即時（JIT）に遅延 IR をコンパイラの内部表現（例：HLO）に変換する。
• 同じ IR グラフが再び出現した場合には再コンパイルを回避するキャッシュを活用し、オーバーヘッドを低減する。
• 非最適化コード（例：制御フロー、動的形状）と最適化されたテンソル演算が共存できるように、明示的なバリア（例：LazyTensorBarrier()）を導入して実行境界を管理する。
• テンソルオブジェクトを、計算が完了するまで解決されないプロンプティブ（Promise）として扱い、非同期実行を可能にする。

実験結果

リサーチクエスチョン

• RQ1 define-by-run ML フレームワークは、ドメイン特化コンパイラによるコンパイルを可能にしつつ、ホスト言語の表現力をフルに活用できるか？
• RQ2 IR グラフ構築のオーバーヘッドを最小限に抑えるにはどうすればよいか？（特に小バッチや低レイテンシ推論ワークロードに適応するため）。
• RQ3 同じ実装を異なるプログラミング言語やハードウェアアクセラレータに跨ってどれだけ再利用できるか？
• RQ4 遅延コンパイルは、実世界の ML ワークロードにおいて、完全トレースやグラフベースのシステムと同等か、それ以上のパフォーマンスを達成できるか？

主な発見

• LazyTensor は、XLA へのコンパイルを可能にしつつ、Python や Swift の機能をフルに活用できることを実証した。これにより、トレーシングベースのシステムに内在する言語サブセット問題を回避できた。
• 同じコア実装が PyTorch と Swift for TensorFlow 間で再利用可能であり、言語やフレームワークを跨ぐ移植性を示した。
• パフォーマンス評価では、HuggingFace Transformers ライブラリを用いた Cloud TPU 上での顕著な高速化が確認され、実世界への適用可能性が裏付けられた。
• 非最適化コードと最適化されたテンソル演算が共存する混合ワークロードを、LazyTensorBarrier() を用いて効果的に管理できることが示された。
• グラフ構築のオーバーヘッドは、大多数の学習ワークロードにおいて許容可能であるが、IR 表現の最適化とキャッシュ化によりさらなる低減が可能である。