著者: シェン・ウェンティン、ファン・イートン、アイ・バオレ、ワン・アン、リー・ヨン

1. はじめに

PAI-TorchAcc (Torch Accelerator) は、Alibaba Cloud AI Platform によって開発された、PyTorch 用の大規模モデル トレーニング アクセラレーション フレームワークです。

PAI-TorchAcc は、PyTorch をベースにしたシンプルで使いやすいインターフェースを提供し、モデル変換なしで HuggingFace 上のモデルにシームレスに接続し、さまざまな分散戦略を使用してトレーニングを加速できます。

PAI-TorchAccは、コミュニティのPyTorch/XLAを活用し、LazyTensorテクノロジーを用いてPyTorchコードを静的実行グラフに変換します。この計算グラフに基づき、Alibaba Cloud上の計算リソースを考慮しながら、GPUハードウェア上でのモデル学習に向けた、広範囲にわたるターゲット分散最適化と計算最適化を実行します。

PAI-TorchAccは、シンプルなモデルアクセス方式と計算グラフベースの最適化により、様々な規模の大規模モデルを柔軟にサポートし、様々なハードウェアとの互換性を備えています。PAI-TorchAccは、1Bから175Bまでの一般的な大規模モデルのトレーニングをサポートし、PyTorchネイティブやMegatron-LMと比較してトレーニングスループットが向上しています。例えば、LLaMAシリーズモデルの場合、PyTorchネイティブと比較して140%、Megatron-LMと比較して5%の向上を実現しています。A100では、MFUは70%に達し、8枚から128枚への線形速度向上は15.6倍に達します。

2. 背景と要件

2.1 背景

• 大規模モデルのトレーニング

近年、大規模言語モデルと動画生成モデルが急速に発展しています。膨大なテキスト、画像、動画のデータセットを用いて学習されたこれらのモデルは、様々な自然言語処理、画像生成、動画生成タスクを実行し、強力な理解・生成能力を備えています。コンピューティングリソースと技術の継続的な進歩により、LLaMA、GPT-3、Tongyi Qianwen、Soraといった大規模モデルのパラメータ数は数億、さらには数兆にまで増加しています。これらのモデルは、多くのベンチマークテストで前例のない性能を発揮しています。

しかし、大規模モデルの学習には非常に高いコストがかかる。例えば、Megatron-LMを用いてOPT-175Bモデルを事前学習する場合、2ヶ月間の学習に数千台のA100が必要であり、ハードウェア利用率は約47%（MFU）であった。この期間中、ハードウェア障害により、モデルは数十回のチェックポイントロードと再学習を経た。PyTorch FSDPを用いたLLaMA-2-70Bの微調整にも、16台のA100と約13.5時間の運用が必要である[2]。NVIDIA A100やH100などのハードウェアリソースは高価で入手が困難であり、よりコストパフォーマンスの高い他のハードウェアリソースが徐々に市場に登場している。

さまざまな大規模モデルの事前トレーニング、再トレーニング、および微調整を高速化し、さまざまなハードウェア リソースを最大限に活用してリソース使用率を向上させることは、大規模モデルのトレーニング コストを削減する効果的な方法です。

• メガトロン-LM

NVIDIA Megatron-LMは、大規模なTransformerベースモデルの学習を目的としたPyTorchベースの分散学習フレームワークです。Megatron-LMは、データ並列性、モデル並列性、パイプライン並列性を統合し、GPT-3などの特定のモデルを学習します。しかし、異なる大規模モデルや学習データセットをMegatron-LMに統合することは柔軟性に欠け、チェックポイントやデータ形式の変換が必要になります。さらに、Megatron-LMは一部のモデル演算子に対して手動最適化を実行しますが、異なるモデルの異なる計算モードを扱う際に、これらの最適化を自動的に適用することは困難です。

• ディープスピード

DeepSpeedは、Microsoftが開発したPyTorch上の大規模モデル向けのオープンソース分散学習フレームワークです。ZeROとパイプライン並列処理をサポートし、Megatron-LMと組み合わせて3D並列処理を実行できます。DeepSpeedは、HuggingFaceトランスフォーマーライブラリの学習コンポーネントとなっています。しかし、DeepSpeedのパフォーマンスは比較的低く、Megatron-LMと同様に、さまざまな計算モードへの柔軟な最適化には限界があります。

• PyTorch/XLA

PyTorch/XLAはPyTorchとOpenXLAを組み合わせ、LazyTenorテクノロジーを用いてPyTorchコードを静的実行グラフに変換します。そして、この静的グラフに対して計算グラフの最適化とバックエンドコンパイルの最適化が実行されます。PyTorch/XLAは主にTPUシナリオ向けに最適化されていますが、Transformersモデルで一般的に使用されるFlashAttentionアクセラレーション演算子のサポート不足、torchrunの起動サポート不足、計算通信の重複の少なさ、メモリオーバーヘッドの高さなど、GPU上では依然としていくつかの問題と最適化の余地が残っています。

2.2 要件

このような背景から、多様なPyTorchモデル、特にTransformer型モデルを容易に統合でき、様々なハードウェアと互換性のある大規模分散トレーニングエンジンが必要です。異なる計算モード、さまざまなハードウェアアーキテクチャ、計算/メモリアクセス能力において、自動的に計算を最適化し、Alibaba Cloudハードウェア上で特に高いパフォーマンスを発揮する必要があります。同時に、モデルのサイズが大きいため、シングルカードメモリやGPUメモリでは完全に収容できません。異なるモデルには、マルチカードトレーニングを完了し、線形速度を向上させるために、異なる分散戦略と適切な通信が必要です。

3. PAI-TorchAccコアテクノロジーの特徴

柔軟なモデル統合

• LLaMA シリーズ、Qwen、BaiChuan、ChatGLM、OLMo、Bloom など、1B から 175B までの一般的な大規模モデルのトレーニングをサポートします。
• HuggingFace のモデルとのシームレスな統合。
• PyTorch モデルへのワンクリック アクセスと高速化。

数千億のモデルパラメータ

• すでに 1B から 175B までの大規模モデルのトレーニングをサポートしています。

包括的なトレーニングモード

• Float32、Float16、BFloat16 などの混合精度トレーニングをサポートします。
• PyTorch モデルの事前トレーニング、微調整、再トレーニングをサポートします。

複合分散戦略

• データ並列、テンソル並列、シーケンス並列、完全シャードデータ並列、パイプラインなどの分散戦略とそれらの組み合わせをサポートします。

自動計算最適化とビデオメモリ最適化

• 手動の勾配チェックポイントと自動再マテリアライゼーションを使用して、ピーク時のメモリ使用量を削減します。
• ビデオ メモリを自動的に計画および管理して、ピーク時のビデオ メモリ使用量を減らし、ビデオ メモリの断片化を減らします。
• カーネルを自動的にコンパイルおよび最適化して計算効率を向上させます。
• 最先端の高性能カーネル (SOTA) に自動的に接続します。

複数のハードウェアと互換性あり

• NVIDIA A100/800、H100/800、V100 などと互換性があります。
• Alibaba Cloud 上の Lingjun クラスターのハードウェア リソースと互換性があります。

既存のフレームワークと比較して

4. PAI-TorchAccアーキテクチャ

4.1 全体的なアーキテクチャ

PAI-TorchAcc のアーキテクチャは、上から下に向かって次のレイヤーに分かれています。

• モデル層: コンピューター ビジョン、自然言語処理、音声合成などのディープラーニング モデルの高速トレーニングをサポートします。
• アルゴリズム ライブラリ: HuggingFace Transformers、PAI-EasyNLP、TIMM などのアルゴリズム ライブラリを使用して構築されたモデルをサポートします。
• フロントエンド: フロントエンド言語として PyTorch を使用したモデルトレーニングをサポートします。
• 低下: LazyTensor や Symbolic Trace などの手法を使用して、フロントエンド コードを静的実行グラフに変換します。
• IR : 高レベルのデバイス非依存 IR と低レベルのデバイス依存 IR を含む多層中間表現を使用し、IR の 2 つのレイヤーでそれぞれ計算グラフの最適化とバックエンドのコンパイルの最適化を実行します。
• コンパイラ最適化エンジン：TorchAccのコンパイラ最適化エンジンには、計算グラフ最適化エンジンであるTorchAcc Compilerと、様々なバックエンドコンパイラ最適化エンジンであるBladeDISCおよびOpenXLAが含まれています。2層IRに基づいて、分散コンピューティング、メモリ、通信、計算、演算子スケジューリング、メモリ管理の最適化を実行し、最適化されたデバイスコードを生成します。
• ハードウェア: 最終的なハードウェア関連のデバイス コードは、異なるコンピューティング能力、帯域幅、ビデオ メモリを備えたハードウェア デバイス上で実行されます。

4.2 インターフェース

PAI-TorchAcc は簡潔なインターフェース セットを抽出し、元のモデル コードを変更することなく、あらゆる PyTorch モデルの柔軟な統合と高速化を可能にします。

PAI-TorchAcc を使用してモデル トレーニングを高速化するには、通常、次の 3 つのステップが必要です。

1. torchacc.Config を定義し、加速オプションを指定します。
2. モデルと設定を渡して torchacc.accelerate を呼び出し、加速トレーニングの準備を完了します。
3. torch data_loader は、データの読み込みを高速化するために torchacc.AsyncLoader を使用してカプセル化されます。

モデル = ...
データローダー = ...
+ # たった1行のコードでモデルを高速化します。Configを渡すことで、分散戦略やコンパイル最適化オプションなど、より豊富な高速化機能を設定することもできます。+ model = torchacc.accelerate(model)
+ # データ読み込みの非同期高速化 + dataloader = torchacc.AsyncLoader(dataloader, model.device)
モデル.トレーニング()
ソースの場合、データローダー内のラベル:
... 

4.3 コンパイル最適化

PAI-TorchAcc は、LazyTensor や Symbolic Trace などのテクノロジーを使用してフロントエンドの PyTorch コードを静的実行グラフに変換し、静的グラフ上で自動最適化を実行して、分散ハードウェア デバイス上での効率的な操作を可能にします。

4.4 計算グラフの最適化

最適化は、High-Level IR (StableHLO) に基づいて、Tensor Graph 上で実行されます。

• 分散: パイプライン並列処理、SPMD などは、グラフ分割と通信演算子の挿入によって実現されます。
• メモリの最適化: オペレータ レベルのメモリの有効範囲と再利用の分析、静的スケジューリング戦略、自動再計算、およびメモリ管理の最適化を通じて、メモリのピークと断片化を削減します。
• 計算の最適化: CSE によって計算を簡素化し、大粒度の演算子の融合によってメモリを大量に消費する演算子を最適化して、カーネルの起動とメモリ アクセスを削減し、計算効率を向上させます。また、自動計算グラフのマッチングと書き換えによって Flash Attention などの高性能カーネルを統合します。
• 通信の最適化: 通信演算子のマージ、分割、非同期化、および演算子のスケジュール設定により、通信効率を向上させ、計算と通信の重複を削減します。

4.5 バックエンドコンパイルの最適化

最適化は、LHLO、LLVM IR、さまざまな MLIR 方言を含む低レベル IR に基づいて、バッファー グラフで実行されます。

• 複数のバックエンド: OpenXLA と Alibaba が独自に開発した BladeDISC コンパイル バックエンドの両方をサポートします。
• ダウングレードとコード生成：上位レベルのStableHLOをLHLOおよび各種MLIR方言にダウングレードし、各レベルのダウングレード処理で最適化を行います。最終的にLLVM IRとして表現され、LLVMを通じ​​てハードウェア最適化されたコードが生成されます。
• カスタム コール: 高レベル IR 自動パターン書き換え用に最適化されたカーネルは、カスタム コールによって呼び出されます。

5.実践的な事例とパフォーマンス

PAI-TorchAccはA100で70%のMFUを達成し、複数のGPU（8GPUから128GPUまで15.6倍の高速化）にほぼ線形にスケーリングします。様々なモデルを柔軟にサポートしながら、その性能はMegatron-LMを凌駕します。私たちは、同じハードウェアリソースを使用して、一般的なオープンソースの大規模モデルで性能テストを実施しました。PAI-TorchAccのトレーニングスループットは、ネイティブPyTorchおよびMegatronと比較して向上しています。例えば、LLaMAシリーズのモデルは、ネイティブPyTorchと比較して140%、Megatronと比較して5%の向上を示しています。

以降の記事では、PAI-TorchAcc を OLMo モデルトレーニングに統合する例とその加速効果、および加速の原因の分析など、具体的な実践的なケーススタディを紹介します。

6.まとめと今後の展望

PAI-TorchAccはPyTorchモデルと柔軟に統合し、並列化戦略、メモリ最適化、計算最適化、スケジューリング最適化を通じて、大規模モデル、視覚モデル、音声モデルの学習を高速化できます。PAI-TorchAccは、LLaMA、LLaMA-2、BaiChuan、ChatGLM、QWen、OLMo、Bloomといった一般的な大規模モデルで既に良好な結果を達成しています。今後も、より多くのシナリオをサポートし、より優れた高速化結果を達成するために、以下の方向で最適化を継続していきます。

1. グラフキャプチャの最適化とサブグラフのコンパイル：計算グラフの生成中に認識できない演算子に遭遇した場合、コンパイルは失敗します。グラフキャプチャをさらに最適化し、サブグラフのコンパイル最適化をサポートします。
2. 自動分散レイアウト：PAI-TorchAccは様々な分散戦略を提供しますが、異なるモデルやハードウェアで最適なパフォーマンスを実現するために、どの戦略の組み合わせとグラフをどのように分割するかを決定するには、経験に基づいた手動設定が必要です。PAI-TorchAccは、静的な計算グラフとモデル/ハードウェア特性を活用して、自動分散レイアウトを実行します。
3. AutoGC: 静的計算グラフとモデル、およびハードウェア特性を活用して、チェックポイントを自動的に選択します。
4. Dynamic Shape のパフォーマンス最適化：Dynamic Shape は再コンパイルによるパフォーマンス低下を引き起こします。現在、バケット化によって再コンパイル回数を削減していますが、それでもパディングは依然として多く残っています。Dynamic Shape のサポートにおけるパフォーマンス向上は、より綿密な最適化の方向性を示しています。
5. 自社開発のコンパイル最適化エンジン BladeDISC の最適化。