[論文レビュー] Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

Keno Fischer, Elliot Saba|arXiv (Cornell University)|Oct 23, 2018

Distributed and Parallel Computing Systems参考文献 4被引用数 9

ひとこと要約

この論文は、XLAコンパイラを介してGoogle Cloud TPUに、ニューラルネットワークモデル（例：VGG19）およびその勾配を含む完全なJuliaプログラムを自動的にオフロードするコンパイラベースの手法を提示する。Juliaの静的解析および型推論の能力を活用することで、前方および後方パスを1つの統合されたTPUカーネルにコンパイルし、100枚の画像バッチで0.23秒（CPU 52.4秒対比）の227倍の高速化を達成。コードは1000行未満で、Juliaコンパイラやコアパッケージの変更は不要である。

ABSTRACT

Google's Cloud TPUs are a promising new hardware architecture for machine learning workloads. They have powered many of Google's milestone machine learning achievements in recent years. Google has now made TPUs available for general use on their cloud platform and as of very recently has opened them up further to allow use by non-TensorFlow frontends. We describe a method and implementation for offloading suitable sections of Julia programs to TPUs via this new API and the Google XLA compiler. Our method is able to completely fuse the forward pass of a VGG19 model expressed as a Julia program into a single TPU executable to be offloaded to the device. Our method composes well with existing compiler-based automatic differentiation techniques on Julia code, and we are thus able to also automatically obtain the VGG19 backwards pass and similarly offload it to the TPU. Targeting TPUs using our compiler, we are able to evaluate the VGG19 forward pass on a batch of 100 images in 0.23s which compares favorably to the 52.4s required for the original model on the CPU. Our implementation is less than 1000 lines of Julia, with no TPU specific changes made to the core Julia compiler or any other Julia packages.

研究の動機と目的

• 一般のJuliaプログラム（機械学習モデルを含む）を、Juliaコンパイラやコアパッケージの変更なしにCloud TPUに完全自動コンパイルすること。
• Zygote.jlとの統合によりエンドツーエンドの自動微分をサポートし、前方および後方パスを両方ともTPUにオフロードすること。
• Juliaの高水準な言語機能（多重ディ patch、高階関数、ジェネリックプログラミングなど）が、TPU実行可能なXLA IRに効果的にコンパイル可能であることを示すこと。
• 制御構造（例：ループ、条件分岐）を含む複雑なワークロード（例：学習ループ）を、1つのカーネルに統合してTPU上で効率的に実行すること。
• コンパイラの変更なしに、スケーラブルで拡張可能で使いやすいインターフェースを提供し、トレーシングベースやフレームワーク特化型の手法の制限を克服すること。

提案手法

• Juliaコンパイラの静的解析および型推論機能を活用して、実行時トレーシングを回避し、Juliaコードを直接XLA中間表現（IR）にコンパイルする。
• Juliaの多重ディ patchおよびジェネリックプログラミングを活用して、Julia IRからXLA IRへの変換を表現・最適化し、型に依存しないコンパイルを実現する。
• データ依存関係を分析し、静的最適化を適用することで、制御構造（例：ループ、条件分岐）をXLA演算に統合する。
• Zygote.jlのソースツーソース自動微分を統合し、前方および後方パスを1つの統合されたTPUカーネルとして生成・オフロードする。
• XLAコンパイラの汎用IRを活用してTPUをターゲットとし、TensorFlow以外のワークロードや線形代数が主な演算である非MLアプリケーションにも対応可能である。
• CPUとTPU間のデータ転送を管理するためのinfeed/outfeed操作を採用し、XLAの型制約に対応するため、StructsOfArraysを介したSoAからAoSへの変換をサポートする。

実験結果

リサーチクエスチョン

• RQ1 制御構造や高水準な抽象化を含む完全なJuliaプログラムが、実行時トレーシングを一切必要とせず、TPU実行可能なXLAカーネルに自動的にコンパイル可能か？
• RQ2 Juliaの多重ディ patchおよびジェネリックプログラミング機能をどれだけ活用して、効率的で型に特化したXLAカーネルを生成できるか？
• RQ3 JuliaのコンパイラインfraとZygoteの自動微分を組み合わせることで、ディープラーニングモデルの前方および後方パスをエンドツーエンドでTPUにオフロード可能か？
• RQ4 完全にオフロードされたVGG19モデルの性能はCPU実行と比べてどの程度向上するか？Cloud TPU上で達成可能な高速化率は？
• RQ5 XLAを介して一般のJuliaコードをTPUにコンパイルする際の、デバッグのしやすさ、型推論の失敗、分散実行の制限といった主な課題は何か？

主な発見

• VGG19モデルの完全な前方パスが1つの統合されたTPUカーネルにコンパイルされ、100枚の画像バッチでCPU実行（52.4秒）と比較して0.23秒で227倍の高速化が達成された。
• Zygote.jlの自動微分とXLAの統合機能を活用することで、前方および後方パスに加え、学習ループごとを含めたエンドツーエンドのオフロードが可能となった。
• 実装はJuliaコードで1000行未満であり、コアJuliaコンパイラや外部パッケージの変更は一切不要であった。
• 安定したパフォーマンス測定が可能であり、TPU計算時間（FluXLA TPU）は総実行時間（wall-clock time）よりもはるかに一貫性があった。
• 線形代数演算が主なワークロードである限り、非MLのJuliaコードに対しても一般化可能であり、深層学習を超えた幅広い応用が期待できる。
• 現在の実装には、型推論の失敗のデバッグが困難であること、XLAの分散プリミティブへの対応が限定的であることといった課題があり、今後の改善の鍵となる分野として特定されている。