Amazon SageMaker で分割されたデータの並列処理を使用して、ほぼ線形のスケーリングで巨大なモデルをトレーニングします

優れた精度を追求するために、自然言語処理やコンピューター ビジョンなどの分野のディープ ラーニング モデルは、過去数年間でサイズが大幅に拡大し、多くの場合、数百億から数千億のパラメーターに数えられています。 これらの巨大なモデルのトレーニングは困難であり、複雑な配布戦略が必要です。 データ サイエンティストと機械学習エンジニアは、トレーニング コンピューティングを最適化する最善の方法を常に探していますが、クラスター全体のサイズとともに増加する可能性がある通信オーバーヘッドと格闘しています。

これが、私たちが最近立ち上げた理由です sハード データ並列処理 on アマゾンセージメーカー、新しいメモリ節約分散トレーニング手法 SageMaker モデル並列 (SMP) ライブラリ. シャード化されたデータの並列処理は、極端な規模のモデル専用に構築されており、Amazon を社内で使用しています MiCS これは、パラメーター収集と勾配同期に起因する高価な通信オーバーヘッドを削減することにより、通信規模を最小限に抑えるための科学的取り組みです。 この新機能は、シーケンス長 30 の 2B パラメータ GPT-2048 モデルで 141 TFLOP を達成し、DeepSpeed ZeRO-39.7 と比較して 3% の速度向上を達成しました。 シーケンス長 10 の 2B GPT-512 モデルの場合、この新しい機能は 564 サンプル/秒も達成し、PyTorch の Fully Sharded Data Parallel (FSDP) と比較して 13.9% の速度向上を実現しました。 巨大なモデルのトレーニングでは、スピードアップのすべてのパーセンテージが、コストの節約とチームの生産性の向上につながることを忘れないでください。

このブログ投稿では、最初にシャード データ並列処理の主な差別化要因と、それをいつ使用するかについて詳しく見ていきます。 次に、この新機能を使用して、SageMaker で 30B パラメータの GPT-2 モデルを簡単にトレーニングする方法を学習します。 最後に、パフォーマンスを他のオープン ソース オプションと比較します。特に、39.7 GPU で DeepSpeed ZeRO を最大 256% 上回っています。

シャード データ並列処理のしくみと使用時期

シャード データ並列処理を紹介する前に、その広範な手法ファミリを見てみましょう。 大規模モデルに対する最近の分散トレーニング アプローチは、モデル パラメーター、勾配、およびオプティマイザーの状態がデータ並列ノード間で共有されるパラダイムに移行しました。 特にフレームワークがサポートしていない場合に、デバイス間で分割するレイヤーを選択するという本質的な複雑さを持つパイプライン並列処理とは異なります。 自動モデル分割、このパラダイムは、モデルが単一の GPU に適合しなければならないというデータ並列処理の制約を取り除きながら、データ並列処理の単純さをエレガントに維持します。

このパラダイムに該当する既存のフレームワーク、特に DeepSpeed ZeRO-3 と FairScale からアップストリームされた PyTorch の FSDP では、モデルの状態は全体でシャーディングされます を GPU: 各 GPU のメモリ消費量を削減する戦略ですが、大きな通信オーバーヘッドが発生するという代償を払います。このオーバーヘッドはクラスターのサイズとともに増加するため、スケーラビリティが大幅に低下します。 対照的に、SMP ライブラリ パーティションのシャード データ並列処理は、モデルの状態を スケール対応 モデル状態の各レプリカを内部のみに分割する方法 サブセット GPUの。

を詳しく見てみましょう スケール対応モデルのパーティショニング シャード データ並列の背後にあるコア テクノロジである MiCS で。 この設計の背後にある直感は、数百億のパラメーターを使用してモデルをトレーニングするために、データ並列グループ全体でトレーニング状態を分割する必要がない可能性があるということです。 たとえば、8 個の V100 GPU (それぞれ 32 GB) は、混合精度を使用して Adam オプティマイザーでトレーニングする場合に約 10 GB のメモリを必要とする 200B パラメーター モデルのモデル状態のレプリカを保持するのに十分です。 モデル状態の完全なレプリカを制限することにより、 最小 GPU のサブセットを使用すると、DeepSpeed や PyTorch FSDP と比較して、通信オーバーヘッドの規模を効果的に削減できます。 シャード データ並列は、階層通信や 2 ホップ勾配同期など、MiCS の他の手法も活用します。 詳細については、こちらをご覧ください AWS での巨大なモデル トレーニングのほぼ線形のスケーリング or MiCS: パブリック クラウドで巨大なモデルをトレーニングするためのニアリニア スケーリング.

では、他の分散トレーニング手法ではなく、シャード データを並列に選択するタイミングをどのように判断すればよいでしょうか? 原則として、モデルのパラメーターが 1 億未満で、GPU メモリに収まる場合、 SageMaker データ並列ライブラリ or SageMaker トレーニング コンパイラ で十分です。 大規模な言語モデルまたはコンピューター ビジョン モデルがある場合は、シャード データ並列処理手法を組み合わせてトレーニングすることをお勧めします。 アクティベーションチェックポイント & アクティベーションオフロード 最初に SageMaker モデル並列ライブラリで、次のような他の手法の前に テンソル並列処理 またはパイプラインの並列処理。

分割データ並列処理を使用して Amazon SageMaker で GPT-2 をトレーニングする

では、複雑さをカプセル化する SMP を使用して、シャード データを並列に使用して GPT-2 モデルをトレーニングする方法を学びましょう。 これ 完全なチュートリアル ノート データ処理、トレーニング ジョブの定義と送信、トレーニング ログの監視まで、プロセス全体を順を追って説明します。 以下は、この機能を使用するための主要な手順を強調した簡単な概要です。

1 はじめに

シャード データ並列処理は PyTorch v1.12.0+ で利用でき、FP16 と BF16 の両方で動作します。 SMP ライブラリを使用する最も簡単な方法は、事前に構築された PyTorch 用の AWS Deep Learning Container を使用することです。 ただし、独自の Docker コンテナーを持ち込みたい場合は、以下を参照できます。 SageMaker Distributed Model Parallel Library を使用して独自の Docker コンテナを作成します。 始めるには、フォローしてください PyTorchトレーニングスクリプトを変更する トレーニング スクリプトで SMP の API を調整します。 このセクションでは、すぐに使用できるトレーニング スクリプトのコード スニペットを使用して、いくつかの主要な手順のみを呼び出します。 train_gpt_simple.py. スクリプト内のコメントをたどることができます。 APIドキュメント SMP API が使用されている場所について詳しく知るには、

まず、呼び出してライブラリをインポートして初期化します smdistributed.modelparallel.torch.init() トレーニング スクリプトの冒頭で:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

次に、分割するモデルをラップします smdistributed.modelParallel.torch.DistributedModel 返されたを使用します DistributedModel 今後のオブジェクト:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

オプティマイザーをラップする smdistributed.modelparallel.torch.DistributedOptimizer オプティマイザの状態を保存およびロードするため。

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

ステップ関数に順方向および逆方向のロジックを配置し、次のように装飾します。 smdistributed.modelparallel.torch.step.  内部で定義された計算 smp.step-decorated 関数は分散的に実行されます。

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. データセットを準備する

私たちは、使用  オープンウェブテキスト この例で使用するデータセットです。 ノートブックはスクリプトを使用します data_prep_512.py データセットをダウンロードして前処理します。 変更することにより、他のデータセットでトレーニングすることもできます data_pipeline.py. 大規模なデータセットとモデルを扱う場合、に保存されているデータを使用してトレーニング ジョブを高速化できます。 光沢のためのAmazonFSxとネイティブに統合された高性能ファイル システムを提供します。 Amazon シンプル ストレージ サービス (S3)。 からの指示を参照してください。 Amazon FSx for Lustre を使用するようにデータ入力チャネルを設定する FSx Lustre ファイル システムをデータ入力チャネルとして設定するためのガイダンスについては、 を参照してください。

3. トレーニング ジョブを開始する

この手順では、すでに トレーニング スクリプトを変更しました 前のセクションで説明したように、データセットを準備しました。 に シャード データの並列処理を有効にするには、 sharded_data_parallel_degree セクションに PyTorch エスティメーター. このチュートリアルでは、 sharded_data_parallel_degree=128 & instace_count=32 p4d.24xlarge ノードの場合、モデルの状態が合計 128 個の GPU のうち 256 個の GPU でシャーディングされることを示します。 この選択された値に基づいて、SMP は自動的にデータ並列度を 2 に設定します (256/128=2 のため)。これは、データ並列処理のために XNUMX つのレプリカがあることを意味します。 の理想的な値を選択するための一般的なルール sharded_data_parallel_degree モデル パラメータ 3B ごとに共有グループにノードを 30 つ追加することです。 このチュートリアルでは、モデル サイズが 10B であるため、シャーディングには少なくとも 16 ノードを使用する必要があります。 また、128 ノード (2 GPU) がしきい値を超える最小の XNUMX のべき乗であるため、 sharded_data_parallel_degree=128.

チェックポインティングのために、チェックポインティング ユーティリティのセットも提供しています。 sharded_data_parallel_checkpoint.py 、完全に再構築するためのユーティリティを含む state_dict 高度なユースケース向け。 最後に、Estimator で fit() を呼び出して、分散トレーニング ジョブを起動できます。

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. トレーニング ジョブを監視する

トレーニング ログにアクセスして、GPU とメモリの使用率を追跡できます。 アマゾンクラウドウォッチ. 「algo-1」のログを確認してください これは、出力ストリームにすべてのインスタンスからのトレーニング ジョブ ログが含まれるメイン ノードであるためです。

ベンチマークパフォーマンス

16 および 32 の p4d.24xlarge ノードの SMP ライブラリで、それぞれシーケンス長 512 および 2048 のシャード データ並列処理のベンチマークを行いました。 30B パラメータの GPT2 モデルは、隠し幅 7168、48 層、および 64 ヘッドを使用するように構成されています。 設定により、シーケンス長が2048のまったく同じ構成を採用できます model_config = "gpt2-30b" チュートリアルノートに。 この設定により、SMP は 73.52 秒あたり 39.7 サンプルを達成し、DeepSpeed ZeRO-3 と比較して 500% の速度向上を達成しました。 トークン サイズが 367 億の場合、この高速化は p4d.24xlarge ノードで約 12,000 時間の節約を意味し、トレーニングごとに XNUMX ドル以上の予算を節約できます! 次の表は、ベンチマークの結果をまとめたものです。

1/ モデル構成ごとに、DeepSpeed ZeRO でさまざまな機能、ステージ、および構成をテストし、DeepSpeed ベースラインとして最高のスループットを提供するものを選択しました。 ベンチマークが実行された アマゾン エラスティック コンピューティング クラウド （Amazon EC2）。 2/ これらの結果は、AWS 向けに最適化された改善されたコミュニケーション コレクティブに依存しており、これはまもなく利用可能になります。 3/ トレーニング時間は、処理されたトークンの数に基づく速度から予測されます。

要約すると、さまざまなモデルと構成で DeepSpeed と比較した場合、SMP のシャード データ並列処理で一貫して高いスループットが観察されました。 この新機能は、DeepSpeed と比較してより優れたメモリ効率も実証し、SMP をより大きなバッチ サイズに適合させ、特定のグローバル バッチ サイズに適合させるために必要な勾配蓄積のレベルを削減できるようにしました。

まとめ

この投稿では、新しい分散トレーニング手法であるシャードデータ並列処理を紹介し、Amazon SageMaker でほぼ線形スケーリングを使用して巨大なモデルのトレーニングを高速化する方法を紹介しました。 また、これに続いて新しい手法を使用して GPT-2 モデルをトレーニングする方法についても説明しました 完全な例。 あなたはフォローすることができます Amazon SageMaker の例 GitHub リポジトリ すべての SageMaker モデルの並列例を追跡するか、次の会議に参加する 分散型トレーニング ワークショップ. シャードされたデータの並列処理の詳細については、 ドキュメント.

著者について

エミリー・ウェバー SageMakerがリリースされた直後にAWSに参加し、それ以来、世界中にそのことを伝えようとしています。 エミリーは、顧客のために新しいML体験を構築する以外に、チベット仏教の瞑想と勉強を楽しんでいます。

缶からく AWS の上級応用科学者であり、AWS で大規模な分散型深層学習を最適化しています。 彼の研究対象は、深層学習、分散最適化、分散システム、および情報理論です。 仕事以外では、サイクリング、旅行、読書、学習を楽しんでいます。

ラフル・ウィゴル AWS のシニア ソフトウェア エンジニアです。 彼は分散型ディープ ラーニング システムに取り組んでおり、クラウドで大規模なディープ ラーニング モデルを簡単かつ効率的にトレーニングできるようにしています。 余暇には、写真、サイクリング、ガーデニングを楽しんでいます。

スヒット・コドゥグル 深層学習フレームワークに取り組んでいる AWS 人工知能グループのソフトウェア開発エンジニアです。 余暇には、ハイキング、旅行、料理を楽しんでいます。

エリン・ホー AWS Deep Learning のプロダクトマネージャーです。 彼女は、お客様が AWS で深層学習モデルを簡単にトレーニングできるようにする製品に取り組んでいます。 仕事以外の楽しみとして、彼女はハイキングとスキーを楽しんでいます。