[論文レビュー] Tensor Program Optimization with Probabilistic Programs

Junru Shao, Xiyou Zhou+8 more|arXiv (Cornell University)|May 26, 2022

Parallel Computing and Optimization Techniques被引用数 10

ひとこと要約

MetaScheduleは、テンソルプログラム最適化のための包括的な探索空間をモジュラーに構築できる分野特化型の確率的プログラミング言語を導入し、ドメインスペシャリストが最適化手法を簡単に拡張できるようにする。ハードウェア固有の変換モジュールと学習駆動型探索を組み合わせることで、TVMよりもエンドツーエンドのディープラーニングワークロードで48%高速化を達成した。

ABSTRACT

Automatic optimization for tensor programs becomes increasingly important as we deploy deep learning in various environments, and efficient optimization relies on a rich search space and effective search. Most existing efforts adopt a search space which lacks the ability to efficiently enable domain experts to grow the search space. This paper introduces MetaSchedule, a domain-specific probabilistic programming language abstraction to construct a rich search space of tensor programs. Our abstraction allows domain experts to analyze the program, and easily propose stochastic choices in a modular way to compose program transformation accordingly. We also build an end-to-end learning-driven framework to find an optimized program for a given search space. Experimental results show that MetaSchedule can cover the search space used in the state-of-the-art tensor program optimization frameworks in a modular way. Additionally, it empowers domain experts to conveniently grow the search space and modularly enhance the system, which brings 48% speedup on end-to-end deep learning workloads.

研究の動機と目的

• 自動テンソルプログラム最適化における探索空間の構築と探索アルゴリズムの分離を図ること。
• コアフレームワークを変更せずにドメインスペシャリストが段階的に探索空間を拡張・カスタマイズできるようにすること。
• 構築された探索空間を用いて最適化されたテンソルプログラムを発見する学習駆動型フレームワークの開発。
• CPUやGPUを含む多様なハードウェアバックエンドで、最先端のシステム（TVM や PyTorch）と同等または優れたパフォーマンスを達成すること。
• モジュラーかつ合成可能な探索空間の構築が、実世界のディープラーニングワークロードにおいて測定可能なパフォーマンス向上をもたらすことを示すこと。

提案手法

• プログラム変換における確率的選択をモデル化する確率的プログラムとして探索空間を表現すること。
• ドメインスペシャリストがモジュラーな変換ルールを確率的プログラムとして定義し、それらを組み合わせられることを可能にすること。
• 微分可能な探索アルゴリズムを用いて確率的探索空間を探索し、低遅延のテンソルプログラムを発見すること。
• 複数の変換モジュール（例：ループタイリング、ベクタライゼーション、テンソル化）を統合された探索空間に組み合わせること。
• GPUテンソルコアなどの専用ハードウェアプリミティブを活用するため、ハードウェア固有のモジュール（例：Use-Tensor-Core）を統合すること。
• フレームワークレベルの変更を最小限に抑えつつ、CPUおよびGPUプラットフォームをカバーするエンドツーエンドの自動チューニングシステムを構築すること。

実験結果

リサーチクエスチョン

• RQ1 確率的プログラミングの抽象化は、テンソルプログラム最適化における探索空間の構築と探索の分離を効果的に実現できるか？
• RQ2 提案された抽象化は、TVM や PyTorch などの最先端フレームワークで使われる最適化手法をどれほど包括的に表現できるか？
• RQ3 変換ルールのモジュラーな合成は、最適化パフォーマンスにどの程度の向上をもたらすか？
• RQ4 エンドツーエンドのディープラーニングモデルにおいて、既存のシステムと同等または優れたパフォーマンスを達成できるか？
• RQ5 ドメインスペシャリストがハードウェア固有の最適化をどの程度効率的に探索空間に拡張できるか？

主な発見

• MetaScheduleは、CPUおよびGPUの両方で、多様な演算子においてTVM（Ansor）と同等またはそれ以上のパフォーマンスを達成した。
• エンドツーエンドのディープラーニングモデルにおいて、MetaScheduleはTVM や PyTorch のパフォーマンスを同等または上回り、ハードウェア固有モジュールを用いることでBERT-Largeで48%の高速化を達成した。
• 2日間で作成した82行のハードウェア固有モジュールが、TVM（AutoTVM）に対して48%のパフォーマンス向上をもたらした。これはカスタマイズの高い実用性を示している。
• 変換モジュールを段階的に組み合わせることで一貫したパフォーマンス向上が得られ、モularityがスケーラブルな探索空間の拡大を可能にしている。
• MetaScheduleの探索空間は、既存のフレームワークの最適化領域を十分にカバーしており、下位層のシステムを細かく修正することなく実現可能である。
• フレームワークはドメインスペシャリストと学習ベースの探索の間で効率的な協働を可能にし、新しい最適化手法を組み込む作業を著しく軽減した。