[論文レビュー] Automap: Towards Ergonomic Automated Parallelism for ML Models

Michael Schaarschmidt, Dominik Grewe+9 more|arXiv (Cornell University)|Dec 6, 2021

Parallel Computing and Optimization Techniques被引用数 5

ひとこと要約

Automapは、誘導的バイアス、探索、学習を統合することで、トランスフォーマーのメガトロン型シャーディング戦略を回復する、プロトタイプの自動パーティショナである。MLモデルにおけるエレガントでコンパイラー統合型のSPMDスタイル並列化を可能にし、プラットフォーム非依存のIRと段階的書き換えを用いることで、数分でエキスパート並のパフォーマンスを達成する。手動での作業を大幅に削減し、最小限のユーザー入力で複雑なモデルにもスケーリング可能である。

ABSTRACT

The rapid rise in demand for training large neural network architectures has brought into focus the need for partitioning strategies, for example by using data, model, or pipeline parallelism. Implementing these methods is increasingly supported through program primitives, but identifying efficient partitioning strategies requires expensive experimentation and expertise. We present the prototype of an automated partitioner that seamlessly integrates into existing compilers and existing user workflows. Our partitioner enables SPMD-style parallelism that encompasses data parallelism and parameter/activation sharding. Through a combination of inductive tactics and search in a platform-independent partitioning IR, automap can recover expert partitioning strategies such as Megatron sharding for transformer layers.

研究の動機と目的

• 大規模なMLモデルにおける効率的なデータ並列、モデル並列、パイプライン並列戦略の自動発見を実現し、エキスパートの知識や手動での書き換えを一切不要とする。
• XLAなどの既存のコンパイラースタック（例：JAX）にシームレスに統合され、特定のハードウェア向けに再実装する必要がない。
• 迅速なインタラクティブなパーティショニングにより、最小限のユーザー入力で実験の時間とコストを削減し、研究の反復を迅速化する。
• 繰り返し構造（例：繰り返しブロック）を活用することで、数万の演算を含む複雑で非構造的なモデルにもスケーリング可能である。
• アーキテクチャの再実装なしに、古いハードウェアを含む多様なアクセラレータファームに実用的なデプロイを可能にする。

提案手法

• 本システムは、XLA HLOの上にレイヤーされた、MLIRにおけるダイアレクトとしての新規でプラットフォーム非依存のパーティショニングIRを用い、シャーディングおよびレプリケーションの意思決定をリライトルールとして表現する。
• モンテカルロツリー探索（MCTS）と学習モデルを組み合わせたハイブリッドアプローチを採用し、高品質な戦略を発見するための試行回数を削減する。
• エキスパートのシャーディングパターン（例：Megatron）を模倣することで誘導的バイアスを組み込み、探索空間を制限し収束を向上させる。
• 段階的書き換えをサポートし、特に共有定数やレイヤーグループを通じて、計算グラフ全体にわたってシャーディング意思決定を伝搬する。
• ユーザーは名前付きスコープ（例：'attention-block'）を介して高レベルのヒントを提供でき、繰り返しレイヤーをグループ化することで、探索の効率性と信頼性を著しく向上できる。
• プロトタイプはJAXで実装されたインタラクションネットワークを用いて合成データで訓練され、Optaxによる最適化とグラフニューラルネットワークによる推論を用いてパーティショニング意思決定を実行する。

実験結果

リサーチクエスチョン

• RQ1 エキスパート最適化戦略（例：Megatron）と同等のパフォーマンスを維持しつつ、最小限のユーザー入力で自動パーティショニングを実現できるか？
• RQ2 数万の演算を含む大規模で非構造的なモデルにおいて、探索の効率をどのように向上できるか？
• RQ3 エキスパートの実践から得た誘導的バイアスは、探索空間をどれほど小さくし、収束をどれほど加速できるか？
• RQ4 コンパイラーのヒント（例：名前付きレイヤーグループ）は、共有定数を介した依存関係伝搬に代わって、深層ネットワークにおける耐障害性を向上させられるか？
• RQ5 学習と探索の組み合わせは、実世界の研究ワークフローにおけるソリューション到達時間にどのように影響するか？

主な発見

• 500回未満の探索エピソード（約数分）で、TPU v3上でMegatron型シャーディングを発見し、Megatronとほぼ同等のパフォーマンスを達成したが、わずかな実行時間のペナルティを伴った。
• MCTSをガイドするための学習済みフィルタを用いることで、近似的に最適な解を信頼性高く発見するためのエピソード数を削減でき、データ駆動型探索の価値を示した。
• コンパイラーのヒント（例：名前付きレイヤーグループ）を用いることで、共有定数の伝搬に依存せずとも、Megatronシャーディングを少数のエピソードで信頼性高く発見できた。
• グループ化や伝搬なしでは、24層のトランスフォーマーにおいてMegatronシャーディングが発見されなかったことから、構造的ヒントがスケーラビリティにおいて果たす重要な役割が明確になった。
• GraphNetsのような非標準的なモデルに対しても、プロトタイプは正常にパーティショニングを実行し、大規模な実験に適した戦略（例：入力エッジシャーディング）を発見した。
• 5万～10万の演算を含むモデルに対してもスケーリング可能であり、インタラクティブな研究ワークフローに適した探索時間（数分、数時間ではない）を達成した。