[論文レビュー] FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads

Zhen Zheng, Pengzhan Zhao+7 more|arXiv (Cornell University)|Sep 23, 2020

Advanced Neural Network Applications参考文献 46被引用数 12

ひとこと要約

FusionStitching は JIT 深層学習コンパイラであり、複雑なデータ依存関係と非一様な並列性を持つ演算子を、大規模な GPU カーネルに統合することで、メモリ集約的深層学習ワークロードを最適化する。共有メモリおよびレジスタシャッフルを用いて中間値の再利用を可能にし、オフチップメモリアクセスとコンテキストスイッチのオーバーヘッドを削減する。最新の技術と比較して最大 2.21× の高速化を達成し、平均で 1.45× の向上を実現した。

ABSTRACT

We show in this work that memory intensive computations can result in severe performance problems due to off-chip memory access and CPU-GPU context switch overheads in a wide range of deep learning models. For this problem, current just-in-time (JIT) kernel fusion and code generation techniques have limitations, such as rough fusion plan exploration strategies and limited code generation ability. We propose FusionStitching, a deep learning compiler capable of fusing memory intensive operators, with varied data dependencies and non-homogeneous parallelism, into large GPU kernels to reduce global memory access and context switch overhead automatically. FusionStitching widens the range of operation combinations that fusion can target beyond previous JIT works by introducing data reuse of intermediate values. It explores large fusion spaces to decide optimal fusion plans with considerations of memory access costs, kernel calls and resource usage constraints. FusionStitching tunes the optimal stitching scheme with a domain-specific cost model efficiently. Experimental results show that FusionStitching can reach up to 2.21x speedup compared to state-of-the-art, with 1.45x on average. Besides these experimental results, we integrated our approach into a compiler product and deployed it onto a production cluster for AI workloads with thousands of GPUs. The system has been in operation for more than 4 months and saves 7,000 GPU hours on average for approximately 30,000 tasks per month.

研究の動機と目的

• オフチップメモリアクセスと CPU-GPU コンテキストスイッチのパフォーマンスボトルネックを解消すること。
• 既存の JIT 統合技術の限界、例えば狭い統合空間とデータ再利用のサポート不足を克服すること。
• 異なるデータ依存関係と並列性を持つ、多様で複雑なメモリ集約的演算子（例：要素演算、転置、還元）の効率的統合を可能にすること。
• 最適な統合計画選択と高性能カーネル生成のためのスケーラブルでドメイン特化されたコストモデルとチューニングフレームワークを開発すること。
• 生産環境の AI クラスタ環境に統合・検証し、実世界のパフォーマンスとリソース節約の有効性を実証すること。

提案手法

• メモリ集約的演算子に複雑な依存関係と非一様な並列性を伴う多様な統合パターンをモデル化するための 4 つのステッチ方式の抽象化を導入する。
• オンチップメモリ（共有メモリ）とレジスタシャッフルを用いてスレッド間での中間値の再利用を可能にし、再計算のオーバーヘッドを低減する。
• メモリアクセスコスト、カーネル起動コスト、リソース使用量を推定することで、統合空間の探索を支援する二段階のドメイン特化コストモデルを設計する。
• 大規模な統合空間内でコスト駆動の知的探索を実行する統合エクスプローラーを実装し、グリーディな局所最適解を回避する。
• スケジューリングチューニングを組み込んだ最適化された GPU カーネルへの統合計画マッピングを実装するコードジェネレータを開発する。
• 動的形状のサポートと本番環境デプロイメントを可能にするために、XLA フレームワークに FusionStitching を統合する。

実験結果

リサーチクエスチョン

• RQ1 JIT 統合における中間データ再利用は、どのように体系的かつ効果的に活用できるか？
• RQ2 複雑で動的なワークロードにおいて、メモリアクセス、カーネル起動オーバーヘッド、リソース使用量の最適なトレードオフを実現する統合パターンとコストモデルは何か？
• RQ3 高度な探索とコード生成を備えた JIT 統合フレームワークは、ルールベースまたはグリーディなアプローチを上回る性能を実現できるか？
• RQ4 中間値の再利用とオンチップメモリの活用は、オフチップメモリトラフィックとコンテキストスイッチのオーバーヘッドをどの程度低減できるか？
• RQ5 動的形状と多数の演算子を有する生産規模の AI ワークロードにおいて、提案された統合フレームワークはどの程度有効であるか？

主な発見

• FusionStitching は、実世界の深層学習モデルにおいて、最新の JIT 最適化技術と比較して最大 2.21× の高速化と平均 1.45× の向上を達成した。
• 多様なメモリ集約的演算子を単一で高度に最適化されたカーネルに統合することで、GPU カーネル起動オーバーヘッドとオフチップメモリアクセスを削減した。
• 共有メモリとレジスタシャッフルによる中間値の再利用を可能にすることで、XLA のような既存の JIT ツールが抱える重複計算の問題を回避した。
• フレームワークは生産用 XLA ベースのコンパイラスタックに成功裏に統合され、数千の GPU を有するクラスタで 4 ヶ月以上にわたり運用された。
• 導入により、30,000 件のタスクで月間約 7,000 GPU 時間のリソース節約が達成され、顕著な実世界の効率向上が実証された。
• ドメイン特化コストモデルにより、動的形状や複雑な演算子の組み合わせに対しても、低チューニングオーバーヘッドで効果的な統合計画選択が可能となった。