[論文レビュー] Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients

William S. Moses, Valentin Churavy|arXiv (Cornell University)|Oct 4, 2020

Parallel Computing and Optimization Techniques参考文献 49被引用数 10

ひとこと要約

この論文では、最適化済みのLLVM中間言語（IR）から直接高速な勾配を合成する、LLVM向けの高パフォーマンスな自動微分（AD）コンパイラプラグインであるEnzymeを紹介する。これにより、外部コードの再実装が不要な、効率的で言語に依存しないADが可能になる。コンパイラ最適化の後にADを実行することで、未最適化IRにおける従来のADと比較して4.5倍の幾何平均の高速化を達成し、機械学習ベンチマークにおいて最先端のツールを上回る性能を発揮する。

ABSTRACT

Applying differentiable programming techniques and machine learning algorithms to foreign programs requires developers to either rewrite their code in a machine learning framework, or otherwise provide derivatives of the foreign code. This paper presents Enzyme, a high-performance automatic differentiation (AD) compiler plugin for the LLVM compiler framework capable of synthesizing gradients of statically analyzable programs expressed in the LLVM intermediate representation (IR). Enzyme synthesizes gradients for programs written in any language whose compiler targets LLVM IR including C, C++, Fortran, Julia, Rust, Swift, MLIR, etc., thereby providing native AD capabilities in these languages. Unlike traditional source-to-source and operator-overloading tools, Enzyme performs AD on optimized IR. On a machine-learning focused benchmark suite including Microsoft's ADBench, AD on optimized IR achieves a geometric mean speedup of 4.5x over AD on IR before optimization allowing Enzyme to achieve state-of-the-art performance. Packaging Enzyme for PyTorch and TensorFlow provides convenient access to gradients of foreign code with state-of-the art performance, enabling foreign code to be directly incorporated into existing machine learning workflows.

研究の動機と目的

• 機械学習フレームワークにおいて外部コードを再実装する必要をなくし、ネイティブで高パフォーマンスな自動微分を可能にすること。
• コンパイラ最適化の前にADを適用することで生じる性能ボトルネックを解消し、非効率な勾配コードを回避すること。
• ソースコードにアクセスできない事前コンパイル済みまたはサードパーティライブラリの間での言語・フレームワークを越えた微分を可能にすること。
• 低レベルで最適化されたIR上でも効率的なADが可能であることを示し、ADは高レベル言語でのみ有効であるという仮定に挑戦すること。
• 科学計算やシミュレーションコードを現代の機械学習ワークフローにスムーズに統合するための道筋を提供すること。

提案手法

• EnzymeはLLVMコンパイラフレームワーク内でのプラグインとして動作し、最適化済みLLVM IR上で直接リバースモード自動微分を実行する。
• ADを適用する前に、LLVMのフルスタック最適化（ループ不変コード移動を含む）を活用することで、パフォーマンス向上を維持する。
• 最適化済みIRのデータフローと制御フローを分析して勾配コードを合成し、効率的なアドジョイントコードを生成する。
• 勾配計算に再帰的テープ構造を採用することで、スタックベースのアプローチと比較してメモリオーバーヘッドを低減する。
• C、C++、Fortran、Julia、Rust、Swiftなど、複数の言語を統一されたIRベースの微分パイプラインで処理可能にし、LLVM IRをターゲットとする。
• PyTorchおよびTensorFlowと統合され、既存のMLワークフロー内で外部コードを直接微分可能にしている。

実験結果

リサーチクエスチョン

• RQ1 未最適化のソースまたはIRではなく、最適化済みLLVM IR上で自動微分を効率的に行うことは可能か？
• RQ2 最適化後にADを適用することで、未最適化コードにおける従来のADと比較して、著しく高速な勾配計算が達成できるか？
• RQ3 コンパイラプラグインベースのADシステムが、多様なMLおよび科学計算ワークロードにおいて最先端のパフォーマンスを達成できるか？
• RQ4 Enzymeは、物理シミュレーターやゲームエンジンなどの外部コードをソース変更なしに機械学習パイプラインにシームレスに統合できるか、その程度はどの程度か？
• RQ5 IRレベルで動作し、GPUコード生成などの高度な機能をサポートする、高パフォーマンスで多言語対応のADシステムを構築することは可能か？

主な発見

• Enzymeは、MicrosoftのADBenchを含む多様なベンチマークスイートにおいて、未最適化IRにおけるAD（Referenceパイプライン）と比較して4.5倍の幾何平均の高速化を達成した。
• 性能向上の主な要因は、後段最適化ADにより、ループ不変コード移動などの最適化が勾配内での重複計算を排除できたことにある。
• ADBenchにおいて評価された2つの高速なC/C++ ADツールすべてのベンチマークでEnzymeが上回り、最先端のパフォーマンスを示した。
• Enzymeが採用する再帰的テープ構造は、スタックベースのテープと比較してメモリオーバーヘッドを低減し、FFTおよびBrussベンチマークでの優れたパフォーマンスに寄与した。
• Enzymeは、事前コンパイル済みライブラリやヘッダオンリーコードの直接微分が可能であり、これはソースツーソースADツールでは不可能である。
• Enzymeは、EulerやRK4ソルバー、LSTM、GMMsといった複雑な科学的ワークロードを効率的に微分でき、専用ツールと比較して同等または優れたパフォーマンスを発揮した。