モデル並列処理の概要

モデル並列処理は、深層学習モデルを複数のデバイス、インスタンス内またはインスタンス間で分割する分散型トレーニング方法です。この紹介ページでは、モデル並列処理の概要、通常はサイズが非常に大きい DL モデルをトレーニングするときに発生する問題を克服するのに役立つ方法の説明、 SageMaker およびモデル並列ライブラリがモデル並列戦略とメモリ消費を管理するのに役立つものの例を示します。

モデル並列処理とは

深層学習モデル (レイヤーとパラメータ) のサイズを大きくすると、コンピュータービジョンや自然言語処理などの複雑なタスクの精度が向上します。ただし、単一の GPU のメモリに収めることができる最大モデルサイズには制限があります。DL モデルをトレーニングする場合、GPU メモリの制限は次の点でボトルネックになる可能性があります。

• モデルのメモリフットプリントはパラメータの数に比例して大きくなるため、トレーニングできるモデルのサイズが制限されます。

• これにより、トレーニング中の GPU ごとのバッチサイズが制限され、GPU の使用率とトレーニング効率が低下します。

1 つの GPU でのモデルのトレーニングに関連する制限を克服するために、 SageMaker はモデル並列ライブラリを提供し、DL モデルを複数のコンピューティングノードに効率的に分散してトレーニングできるようにします。さらに、このライブラリでは、EFA 対応デバイスを使用して最適化された分散トレーニングを実現できます。これにより、低遅延、高スループット、OS バイパスによるノード間通信のパフォーマンスが向上します。

モデル並列処理を使用する前にメモリ要件を見積もる

SageMaker モデル並列ライブラリを使用する前に、大規模な DL モデルをトレーニングする際のメモリ要件を理解するために、次の点を考慮してください。

AMP (FP16) と Adam オプティマイザを使用するトレーニングジョブでは、パラメータ 1 つあたりに必要な GPU メモリは約 20 バイトで、以下のように分類できます。

• FP16 パラメータ (最大 2 バイト)

• FP16 勾配 (最大 2 バイト)

• FP32 オプティマイザステート (Adam オプティマイザに基づく最大 8 バイト)

• パラメータの FP32 コピー (最大 4 バイト) (optimizer apply (OA) オペレーションに必要)

• 勾配の FP32 コピー (最大 4 バイト) (OA オペレーションに必要)

100 億個のパラメータを持つ比較的小規模な DL モデルでも、少なくとも 200GB のメモリが必要になることがあります。これは、単一の GPU で利用できる一般的な GPU メモリ (例えば、40GB/80GB のメモリを搭載した NVIDIA A100 と 16/32GB のメモリを搭載した V100) よりもはるかに大きいメモリです。モデルとオプティマイザの状態に必要なメモリ量に加えて、フォワードパスで生成されるアクティベーションなど、他のメモリ消費量もあることに注意してください。必要なメモリは 200GB をはるかに超える場合があります。

分散型トレーニングでは、それぞれ NVIDIA V100 と A100 Tensor コア GPU を搭載した Amazon EC2 P3 インスタンスと P4 インスタンスを使用することをお勧めします。CPU コア、RAM、アタッチされたストレージボリューム、ネットワーク帯域幅などの仕様の詳細については、Amazon EC2 インスタンスタイプページの「Accelerated Computing」セクションを参照してください。

高速コンピューティングインスタンスを使用しても、Megatron-LM や T5 のように約 100 億個のパラメータを持つモデルや、GPT-3 のように数千億個のパラメータを持つさらに大規模なモデルでは、各 GPU デバイスのモデルレプリカに適合しないことは明らかです。

ライブラリがモデルの並列処理とメモリ節約の手法を採用している方法

このライブラリは、さまざまなタイプのモデル並列処理機能と、オプティマイザ状態のシャーディング、アクティベーションチェックポイント、アクティベーションオフロードなどのメモリ節約機能で構成されています。これらの手法をすべて組み合わせることで、数千億のパラメータで構成される大規模なモデルを効率的にトレーニングできます。

シャーディングデータ並列処理 ( で使用可能 PyTorch）

シャーディングデータ並列処理は、メモリを節約する分散トレーニング手法で、モデルの状態 (モデルパラメータ、勾配、およびオプティマイザの状態) をデータ並列グループ内の GPU 間で分割します。

SageMaker は、MiCS の実装を通じてシャーディングデータ並列処理を実装します。これは c の ommunication s cale を模倣し、 のブログ記事「 のギガネティックモデルトレーニングのほぼ線形スケーリング AWS」で説明されているライブラリです。

シャーディングデータ並列処理は、スタンドアロン戦略としてモデルに適用できます。さらに、NVIDIA A100 Tensor Core GPU、ml.p4d.24xlarge を搭載した最もパフォーマンスの高い GPU インスタンスを使用している場合は、SMDDP Collectives が提供する AllGather オペレーションによるトレーニング速度の向上という利点を活用できます。

シャーディングデータ並列処理について深く掘り下げ、その設定方法や、シャーディングデータ並列処理とテンソル並列処理や FP16 トレーニングなどの他の手法を組み合わせて使用する方法については、シャーディングデータ並列処理 を参照してください。

パイプラインの並列処理 ( PyTorch および で使用可能 TensorFlow）

パイプライン並列処理は、レイヤーのセットまたはオペレーションをデバイスのセット全体に分割し、各オペレーションはそのまま残します。モデルパーティションの数 (pipeline_parallel_degree) の値を指定する場合、GPU の総数 (processes_per_host) はモデルパーティションの数で割り切れる必要があります。これを正しく設定するには、pipeline_parallel_degree および processes_per_host のパラメータに正しい値を指定する必要があります。簡単な計算は以下のようになります。

(pipeline_parallel_degree) x (data_parallel_degree) = processes_per_host

指定した 2 つの入力パラメータを指定すると、ライブラリがモデルレプリカ (別名 data_parallel_degree) の数を計算します。

例えば、ml.p3.16xlarge のような 8 つの GPU ワーカーを持つ ML インスタンスを "pipeline_parallel_degree": 2 および "processes_per_host": 8 に設定して使用すると、ライブラリは GPU 全体に分散モデルと 4 方向データ並列処理を自動的にセットアップします。次の図は、モデルを 8 つの GPU に分散して 4 方向データ並列処理と双方向パイプライン並列処理を実現する方法を示しています。パイプライン並列グループとして定義し、PP_GROUP というラベルを付けた各モデルレプリカは、2 つの GPU に分割されます。モデルの各パーティションは 4 つの GPU に割り当てられ、4 つのパーティションレプリカはデータ並列グループに属し、DP_GROUP というラベルが付けられます。テンソル並列処理がなければ、パイプライン並列グループは基本的にモデル並列グループになります。

パイプラインの並列処理について詳しくは、「 SageMaker モデル並列処理ライブラリのコア機能」を参照してください。

パイプライン並列処理を使用してモデルの実行を開始するには、 SageMaker 「モデル並列ライブラリを使用して SageMaker 分散トレーニングジョブを実行する」を参照してください。

テンソル並列処理 ( で使用可能 PyTorch）

Tensor 並列処理は、個々のレイヤー (nn.Modules) をデバイス間で分割し、並列に実行します。次の図は、ライブラリがモデルを 4 つのレイヤーに分割して双方向のテンソル並列処理 ("tensor_parallel_degree": 2) を実現する方法の最も単純な例を示しています。各モデルレプリカのレイヤーは 2 分割され、2 つの GPU に分散されます。この例の場合、モデル並列設定には "pipeline_parallel_degree": 1と "ddp": True (バックグラウンドで PyTorch DistributedDataParallel パッケージを使用) も含まれるため、データ並列処理の度合いは 8 になります。このライブラリは、テンソル分散モデルのレプリカ間の通信を管理します。

この機能の便利さは、特定のレイヤーまたはレイヤーのサブセットを選択してテンソル並列処理を適用できる点にあります。のテンソル並列処理やその他のメモリ節約機能の詳細と PyTorch、パイプラインとテンソル並列処理の組み合わせを設定する方法については、「」を参照してくださいテンソル並列処理。

オプティマイザ状態のシャーディング ( で使用可能 PyTorch）

ライブラリがオプティマイザステートシャーディングをどのように実行するのかを理解するために、4 つのレイヤーからなる単純なモデル例を考えてみましょう。ステートシャーディングを最適化するうえで重要なのは、オプティマイザの状態をすべての GPU に複製する必要がないということです。代わりに、オプティマイザの状態の単一のレプリカが、デバイス間の冗長性なくデータ並列ランクにわたってシャーディングされます。例えば、GPU 0 はレイヤー 1 のオプティマイザ状態を保持し、次の GPU 1 は L2 のオプティマイザ状態を保持する、という具合です。次のアニメーション図は、オプティマイザ状態のシャーディング手法による後方伝播を示しています。後方伝播の最後には、optimizer apply (OA) オペレーションでオプティマイザの状態を更新し、all-gather (AG) オペレーションで次の反復に備えてモデルパラメータを更新するための計算時間とネットワーク時間が残ります。最も重要なのは、この reduce 処理が GPU 0 での計算と重複する可能性があるため、メモリ効率が高まり、後方伝播が速くなることです。現在の実装では、AG と OA のオペレーションは compute と重複しません。その結果、AG オペレーション中の計算時間が長くなる可能性があるため、トレードオフが発生する可能性があります。

この機能の使用方法の詳細については、「Optimizer State Sharding」を参照してください。

アクティベーションのオフロードとチェックポイント ( で利用可能 PyTorch）

GPU メモリを節約するため、ライブラリはアクティベーションチェックポイントをサポートしています。これにより、フォワードパス中にユーザー指定モジュールの内部アクティベーションが GPU メモリに保存されるのを防ぐことができます。ライブラリはバックワードパス中にこれらのアクティベーションを再計算します。さらに、アクティベーションオフロード機能は、保存されているアクティベーションを CPU メモリにオフロードし、バックワードパス中に GPU にフェッチバックすることで、アクティベーションメモリのフットプリントをさらに削減します。これらの機能の使用方法について詳しくは、「Activation Checkpointing」と「Activation Offloading」を参照してください。

モデルに適した手法の選択

適切な手法と設定の選択の詳細については、SageMaker 「分散モデル並列のベストプラクティス」と「設定のヒントと落とし穴」を参照してください。