[論文レビュー] Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models

Yuliang Liu, Shenggui Li+4 more|arXiv (Cornell University)|Feb 6, 2023

Parallel Computing and Optimization Techniques被引用数 6

ひとこと要約

Colossal-Auto は、大規模モデルのマルチGPUシステムにおけるトレーニングにおいて、内部演算子並列化とアクティベーションチェックポイント化を統合的に最適化する自動化された統一システムを導入する。ハードウェアに適したプロファイリングと二段階ソルバーを組み合わせることで、最小限のコード変更で近似的に最適な分散実行計画を生成し、従来の手法と比較して A100 GPU において優れた弱スケーリング性能を達成する。

ABSTRACT

In recent years, large-scale models have demonstrated state-of-the-art performance across various domains. However, training such models requires various techniques to address the problem of limited computing power and memory on devices such as GPUs. Some commonly used techniques include pipeline parallelism, tensor parallelism, and activation checkpointing. While existing works have focused on finding efficient distributed execution plans (Zheng et al. 2022) and activation checkpoint scheduling (Herrmann et al. 2019, Beaumont et al. 2021}, there has been no method proposed to optimize these two plans jointly. Moreover, ahead-of-time compilation relies heavily on accurate memory and computing overhead estimation, which is often time-consuming and misleading. Existing training systems and machine learning pipelines either physically execute each operand or estimate memory usage with a scaled input tensor. To address these challenges, we introduce a system that can jointly optimize distributed execution and gradient checkpointing plans. Additionally, we provide an easy-to-use symbolic profiler that generates memory and computing statistics for any PyTorch model with a minimal time cost. Our approach allows users to parallelize their model training on the given hardware with minimum code change based. The source code is publicly available at Colossal-AI GitHub or https://github.com/hpcaitech/ColossalAI

研究の動機と目的

• 大規模モデルの分散トレーニング戦略を手動でチューニングする課題に対処すること。これは深い専門知識を要し、誤りを起こしやすい。
• 従来のシステムで別々に扱われる傾向にある、内部演算子並列化とアクティベーションチェックポイント化を統合的に最適化すること。
• 分散トレーニング経験のない研究者や実務家が、大規模モデルを効率的にトレーニングできるように障壁を低減すること。
• 複雑なGPUトポロジー（NVLink や PCIe トランスポータを含む）に適応可能な、自動的かつハードウェアに適した実行計画生成を可能にすること。
• 任意次元のデバイスメッシュと新しいハードウェア構成に対応する柔軟で拡張可能な並列化戦略をサポートすること。

提案手法

• シンボリックプロファイラーは、計算グラフ上でメタ実行を実行し、最小限のランタイムコストで細粒度のメモリおよび計算オーバーヘッドを収集する。
• 二段階ソルバーは、最初に最適な内部演算子並列化戦略を特定し、次にメモリ予算を満たすようにアクティベーションチェックポイントスケジューリングを共同で最適化する。
• ヒューリスティックなテンソルレイアウト変換アルゴリズムにより、任意のN次元デバイスメッシュに効率的なシャーディングが可能となり、複雑なトポロジーをサポートする。
• システムは、演算をチェックポイント付きで並列化された関数に置き換えることで、PyTorch モデルを分散実行コードにコンパイルする。
• 実行計画生成は、NVLink や PCIe の帯域幅の違いを含むハードウェアトポロジーを考慮し、通信と計算を重ね合わせる。
• システムは PyTorch と統合されており、最小限のユーザー側コード変更で自動的にコード再コンパイルを可能にする。

実験結果

リサーチクエスチョン

• RQ1 内部演算子並列化とアクティベーションチェックポイント化を統合的に最適化することで、大規模モデルトレーニングにおけるメモリおよび計算効率を向上させることができるか？
• RQ2 シンボリックプロファイラーは、低コストで PyTorch モデルのメモリおよび計算コストを正確に推定できるか？
• RQ3 統合ソルバーは、NVLink や PCIe などの複雑なGPUインターコネクトトポロジーに適応可能なハードウェアに適した実行計画を生成できるか？
• RQ4 並列化とチェックポイント化の統合最適化は、逐次的または独立した最適化と比較して、トレーニング性能およびメモリ使用量においてどのように異なるか？
• RQ5 手動チューニングなしで、異なるモデルアーキテクチャやハードウェア構成にどの程度一般化できるか？

主な発見

• Colossal-Auto は、全評価手法の中で最も優れた弱スケーリング性能を達成し、14.55Bパラメータのモデルに対して 8 個の A100 GPU で 0.824 PFLOPS を達成した。
• 全実験設定において、Megatron-LM や Optimus、3D タンジェント並列化よりも優れた性能を示し、特に大規模モデルにおいて顕著であった。
• 実行計画は自動的に高帯域幅の NVLink リンクを活用し、低帯域幅の PCIe 通信を計算と重ね合わせる。
• データ並列化（DDP）とは異なり、大規模スケールでもメモリ不足エラーを回避した。
• シンボリックプロファイラーは、計算グラフ全体にわたり正確かつ低コストでメモリおよび計算使用量をプロファイリングでき、効果的な最適化を可能にした。
• 並列化とチェックポイント化の統合最適化は、それぞれのコンponentを個別に最適化するのと比較して、より良いメモリ利用効率とパフォーマンスを実現した。