シャーディングデータ並列処理 - Amazon SageMaker

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

シャーディングデータ並列処理

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

注記

シャーディングデータ並列処理は、 SageMaker モデル並列処理ライブラリ v1.11.0 以降 PyTorch で で使用できます。

トレーニングジョブを大規模なGPUクラスターにスケールアップする場合、モデルのトレーニング状態を複数の にシャーディングすることで、モデルのメモリごとのGPUフットプリントを削減できますGPUs。これにより 2 つの利点が返されます。より大きなモデルを適合させることができ、そうしないと標準のデータ並列処理でメモリが不足するか、解放されたGPUメモリを使用してバッチサイズを増やすことができます。

標準データ並列処理手法は、データ並列グループの 全体でトレーニング状態GPUsをレプリケートし、 AllReduceオペレーションに基づいて勾配集約を実行します。シャーディングデータ並列処理は、オプティマイザ状態のシャーディングの性質を考慮して、標準のデータ並列分散トレーニング手順を変更します。モデル状態とオプティマイザ状態がシャーディングされるランクのグループは、シャーディンググループと呼ばれます。シャーディングデータ並列処理手法は、モデルのトレーニング可能なパラメータと、それに対応する勾配とオプティマイザの状態をシャーディンググループ GPUsの 全体にシャーディングします。

SageMaker は、MiCS の実装を通じてシャーディングデータ並列処理を実現します。これは、 の AWS ブログ記事「巨大なモデルトレーニングのほぼ線形スケーリング」で AWS説明されています。この実装では、シャーディング度を設定可能なパラメータとして設定できます。このパラメータは、データ並列処理度よりも小さくなければなりません。フォワードパスとバックワードパスのたびに、MiCS は AllGatherオペレーションGPUsを通じて一時的にすべての のモデルパラメータをリコンバインします。各レイヤーのフォワードパスまたはバックワードパスの後、MiCS はパラメータを再度シャードしてGPUメモリを節約します。バックワードパス中、MiCS は勾配を減らし、同時に ReduceScatterオペレーションGPUsを通じてそれらをシャードします。最後に、MiCS は、オプティマイザ状態のローカルシャードを使用して、ローカルの縮小勾配とシャーディング勾配を対応するローカルパラメータシャードに適用します。通信オーバーヘッドを減らすために、 SageMaker モデル並列処理ライブラリはフォワードパスまたはバックワードパスで今後のレイヤーをプリフェッチし、ネットワーク通信と計算を重複させます。

モデルのトレーニング状態はシャーディンググループ全体で複製されます。つまり、勾配をパラメータに適用する前に、シャーディンググループ内で行われる ReduceScatter オペレーションに加えて、シャーディンググループ全体で AllReduce オペレーションを実行する必要があります。

実際には、シャーディングデータ並列処理は、通信オーバーヘッドとGPUメモリ効率の間にトレードオフをもたらします。シャーディングデータ並列処理を使用すると通信コストが増加しますが、 あたりのメモリフットプリント GPU (アクティベーションによるメモリ使用量を除く) はシャーディングデータ並列処理度で割られるため、より大きなモデルをGPUクラスターに収めることができます。

シャーディングデータ並列処理度の選択

シャーディングデータ並列処理度の値を選択するとき、その値によってデータ並列処理度を均等に分割する必要があります。例えば、8 方向データ並列処理ジョブの場合、シャーディングデータ並列処理度に 2、4、または 8 を選択します。シャーディングデータ並列処理度を選択するときは、小さい数値から始め、モデルが目的のバッチサイズとともにメモリに収まるまで徐々に増やすことをお勧めします。

バッチサイズの選択

シャーディングデータ並列処理を設定したら、GPUクラスターで正常に実行できる最適なトレーニング設定を見つけてください。大規模言語モデル (LLM) をトレーニングする場合は、バッチサイズ 1 から開始し、 out-of-memory (OOM) エラーを受信する時点に達するまで徐々に増やします。最小のバッチサイズでもOOMエラーが発生した場合は、より高いレベルのシャーディングデータ並列処理を適用するか、シャーディングデータ並列処理とテンソル並列処理の組み合わせを適用します。

シャーディングデータ並列処理をトレーニングジョブに適用する方法

シャーディングデータ並列処理を開始するには、トレーニングスクリプトに必要な変更を適用し、 sharded-data-parallelism-specific パラメータを使用して SageMaker PyTorch 推定器を設定します。また、開始点として、参照値とサンプルノートブックを使用することも検討してください。

PyTorch トレーニングスクリプトを適応させる

「ステップ 1: PyTorch トレーニングスクリプトを変更して、モデルオブジェクトとオプティマイザオブジェクトを torch.nn.parallel および torch.distributedモジュールのsmdistributed.modelparallel.torchラッパーでラップする」の手順に従います。

(オプション) 外部モデルパラメータを登録するための追加変更

モデルが で構築torch.nn.Moduleされ、モジュールクラス内で定義されていないパラメータを使用している場合は、 が 中に完全なパラメータを収集SMPできるようにモジュールに手動で登録する必要があります。パラメータをモジュールに登録するには、「smp.register_parameter(module, parameter)」を使用してください。

class Module(torch.nn.Module): def __init__(self, *args): super().__init__(self, *args) self.layer1 = Layer1() self.layer2 = Layer2() smp.register_parameter(self, self.layer1.weight) def forward(self, input): x = self.layer1(input) # self.layer1.weight is required by self.layer2.forward y = self.layer2(x, self.layer1.weight) return y

SageMaker PyTorch 推定器をセットアップする

で SageMaker PyTorch 推定器を設定するときはステップ 2: SageMaker Python SDK を使用してトレーニングJob を起動する、シャーディングデータ並列処理のパラメータを追加します。

シャーディングデータ並列処理を有効にするには、 sharded_data_parallel_degreeパラメータを推定器に追加します SageMaker PyTorch。このパラメータは、トレーニング状態がシャーディングGPUsされる の数を指定します。sharded_data_parallel_degree の値は、1 とデータ並列処理度の間の整数で、データ並列処理度を均等に割る必要があります。ライブラリは の数を自動的に検出GPUsするため、データ並列度であることに注意してください。シャーディングデータ並列処理には、次の追加パラメータが利用できます。

  • "sdp_reduce_bucket_size" (int、デフォルト: 5e8) — デフォルトの dtype の要素数でPyTorch DDP勾配バケットのサイズを指定します。

  • "sdp_param_persistence_threshold" (int、デフォルト: 1e6) — パラメータテンソルのサイズを、各 で保持できる要素の数で指定しますGPU。シャーディングデータ並列処理は、各パラメータテンソルGPUsをデータ並列グループ間で分割します。パラメータテンソルの要素数がこのしきい値より小さい場合、パラメータテンソルは分割されません。これにより、パラメータテンソルがデータ並列 間でレプリケートされるため、通信オーバーヘッドが軽減されますGPUs。

  • "sdp_max_live_parameters" (int、デフォルト: 1e9) — フォワードパスとバックワードパス中に再結合されたトレーニング状態に同時に配置できるパラメータの最大数を指定します。アクティブなパラメータの数が所定のしきい値に達すると、AllGather オペレーションによるパラメータの取得は一時停止します。このパラメータをの数を増加すると、メモリ使用量が増えることに注意してください。

  • "sdp_hierarchical_allgather" (bool、デフォルト: True)True に設定すると、AllGather オペレーションは階層的に実行されます。最初に各ノード内で実行され、次にノード間で実行されます。複数ノードの分散型トレーニングジョブでは、階層型 AllGather オペレーションが自動的に有効になります。

  • "sdp_gradient_clipping" (float、デフォルト: 1.0) — モデルパラメータを介して逆方向に伝播する前に、勾配の L2 ノルムを勾配クリッピングするしきい値を指定します。シャーディングデータ並列処理を有効にすると、勾配クリッピングも有効になります。デフォルトのしきい値は 1.0 です。勾配が爆発的に増加する問題がある場合は、このパラメータを調整してください。

次のコードには、シャーディングデータ並列処理の例が示されています。

import sagemaker from sagemaker.pytorch import PyTorch smp_options = { "enabled": True, "parameters": { # "pipeline_parallel_degree": 1, # Optional, default is 1 # "tensor_parallel_degree": 1, # Optional, default is 1 "ddp": True, # parameters for sharded data parallelism "sharded_data_parallel_degree": 2, # Add this to activate sharded data parallelism "sdp_reduce_bucket_size": int(5e8), # Optional "sdp_param_persistence_threshold": int(1e6), # Optional "sdp_max_live_parameters": int(1e9), # Optional "sdp_hierarchical_allgather": True, # Optional "sdp_gradient_clipping": 1.0 # Optional } } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8 # Required } smp_estimator = PyTorch( entry_point="your_training_script.py", # Specify your train script role=sagemaker.get_execution_role(), instance_count=1, instance_type='ml.p3.16xlarge', framework_version='1.13.1', py_version='py3', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="sharded-data-parallel-job" ) smp_estimator.fit('s3://my_bucket/my_training_data/')

リファレンス設定

SageMaker 分散トレーニングチームは、開始点として使用できる以下のリファレンス設定を提供します。次の設定から推定して、モデル設定のGPUメモリ使用量を試して推定できます。

SMDDP Collectives によるシャーディングデータ並列処理

モデル/パラメータの数 Num インスタンス インスタンスタイプ シーケンスの長さ グローバルバッチサイズ ミニバッチサイズ シャーディングデータ並列度
GPT-NEOX-20B 2 ml.p4d.24xlarge 2048 64 4 16
GPT-NEOX-20B 8 ml.p4d.24xlarge 2048 768 12 32

例えば、200 億パラメータのモデルのシーケンス長を増やしたり、モデルのサイズを 650 億パラメータに増やす場合は、最初にバッチサイズを減少させる必要があります。それでもモデルが最小のバッチサイズ (バッチサイズ 1) に収まらない場合は、モデルの並列処理度を上げてみます。

テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理

モデル/パラメータの数 Num インスタンス インスタンスタイプ シーケンスの長さ グローバルバッチサイズ ミニバッチサイズ シャーディングデータ並列度 テンソル並列度 アクティベーションオフロード
GPT-NEOX-65B 64 ml.p4d.24xlarge 2048 512 8 16 8 Y
GPT-NEOX-65B 64 ml.p4d.24xlarge 4096 512 2 64 2 Y

シャーディングデータ並列処理とテンソル並列処理を組み合わせて使用すると、シーケンス長が長いテキストデータを使用しながら、大規模な言語モデル (LLM) を大規模なクラスターに適合させ、その結果、バッチサイズを小さくし、GPUメモリ使用量を処理して長いテキストシーケンスLLMsに対してトレーニングする場合に便利です。詳細については、「テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理」を参照してください。

導入事例、ベンチマーク、その他の設定例については、ブログ記事「Amazon SageMaker モデル並列ライブラリ でのパフォーマンス向上」を参照してください。

SMDDP Collectives によるシャーディングデータ並列処理

SageMaker データ並列処理ライブラリは、 AWS インフラストラクチャに最適化された集合通信プリミティブ (SMDDP 集合) を提供します。Elastic Fabric Adapter (EFA) を使用して all-to-all-type 通信パターンを採用することで最適化を実現し、高スループットでレイテンシーの影響を受けにくい集合体を実現し、通信関連の処理を にオフロードしCPU、計算GPUサイクルを解放します。大規模なクラスターでは、 SMDDP Collectives は と比較して分散トレーニングのパフォーマンスを最大 40% 向上させることができますNCCL。導入事例とベンチマークの結果については、ブログ「Amazon SageMaker モデル並列処理ライブラリ」の「新しいパフォーマンスの向上」を参照してください。

注記

SMDDP Collectives によるシャーディングデータ並列処理は、 SageMaker モデル並列処理ライブラリ v1.13.0 以降、および SageMaker データ並列処理ライブラリ v1.6.0 以降で使用できます。Collectives Supported configurationsでシャーディングデータ並列処理を使用するには、SMDDP「」も参照してください。

大規模な分散トレーニングで一般的に使用される手法であるシャーディングデータ並列処理では、 AllGather集合体を使用して、計算と並行してフォワードパス計算とバックワードパス計算用のシャーディングレイヤーパラメータを再構成しますGPU。大規模なモデルでは、GPUボトルネックの問題を回避し、トレーニング速度を低下させるために、AllGatherオペレーションを効率的に実行することが重要です。シャーディングデータ並列処理がアクティブ化されると、 SMDDP Collectives はこれらのパフォーマンスが重要なAllGather集合体にドロップされ、トレーニングスループットが向上します。

Collectives SMDDP でトレーニングする

トレーニングジョブでシャーディングデータ並列処理がアクティブ化され、 を満たすとSupported configurations、 SMDDP 集合体は自動的にアクティブ化されます。内部的には、 Collectives SMDDP は、 AWS インフラストラクチャでパフォーマンスを発揮するようにAllGather集合体を最適化し、他のすべての集合NCCL体については にフォールバックします。さらに、サポートされていない設定では、 を含むすべての集合体がNCCLバックエンドAllGatherを自動的に使用します。

SageMaker モデル並列処理ライブラリバージョン 1.13.0 以降、 "ddp_dist_backend"パラメータがmodelparallelオプションに追加されます。この設定パラメータのデフォルト値は で"auto"、可能な限り SMDDP Collectives を使用し、NCCLそれ以外の場合はフォールバックします。ライブラリが常に を使用するように強制するにはNCCL、 "ddp_dist_backend"を設定パラメータ"nccl"に指定します。

次のコード例は、 "ddp_dist_backend"パラメータでシャーディングデータ並列処理を使用して PyTorch 推定器を設定する方法を示しています。このパラメータは"auto"デフォルトで に設定されるため、追加はオプションです。

import sagemaker from sagemaker.pytorch import PyTorch smp_options = { "enabled":True, "parameters": { "partitions": 1, "ddp": True, "sharded_data_parallel_degree": 64 "bf16": True, "ddp_dist_backend": "auto" # Specify "nccl" to force to use NCCL. } } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8 # Required } smd_mp_estimator = PyTorch( entry_point="your_training_script.py", # Specify your train script source_dir="location_to_your_script", role=sagemaker.get_execution_role(), instance_count=8, instance_type='ml.p4d.24xlarge', framework_version='1.13.1', py_version='py3', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="sharded-data-parallel-demo", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/')

サポートされている設定

SMDDP Collectives による AllGatherオペレーションは、次のすべての設定要件が満たされると、トレーニングジョブでアクティブ化されます。

  • シャーディングデータ並列処理度が 1 を超えている。

  • Instance_count が 1 より大きい

  • Instance_typeml.p4d.24xlarge に等しい

  • SageMaker PyTorch v1.12.1 以降のトレーニングコンテナ

  • SageMaker データ並列処理ライブラリ v1.6.0 以降

  • SageMaker モデル並列処理ライブラリ v1.13.0 以降

パフォーマンスとメモリのチューニング

SMDDP 集合体は追加のGPUメモリを使用します。モデルトレーニングのユースケースに応じてGPUメモリ使用量を設定する環境変数は 2 つあります。

  • SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES – SMDDPAllGatherオペレーション中、AllGather入力バッファはノード間通信用の一時バッファにコピーされます。SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES 変数は、この一時バッファのサイズ (バイト単位) を制御します。一時バッファのサイズがAllGather入力バッファのサイズよりも小さい場合、AllGather集合体は を使用するようにフォールバックしますNCCL。

    • デフォルト値: 16 * 1024 * 1024 (16 MB)

    • 許容値: 8192 の任意の倍数

  • SMDDP_AG_SORT_BUFFER_SIZE_BYTESSMDDP_AG_SORT_BUFFER_SIZE_BYTES 変数は、ノード間通信から収集されたデータを格納するための一時バッファのサイズを (バイト単位で) 設定します。この一時バッファのサイズが より小さい場合1/8 * sharded_data_parallel_degree * AllGather input sizeAllGather集合体は を使用するようにフォールバックしますNCCL。

    • デフォルト値: 128 * 1024 * 1024 (128 MB)

    • 許容値: 8192 の任意の倍数

バッファサイズ変数に関するチューニングガイダンス

環境変数のデフォルト値は、ほとんどのユースケースで適切に機能します。これらの変数は、トレーニングが out-of-memory (OOM) エラーになった場合にのみ調整することをお勧めします。

次のリストでは、 Collectives のパフォーマンスの向上を維持しながら Collectives SMDDP のGPUメモリフットプリントを削減するためのチューニングのヒントについて説明します。

  • SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES のチューニング

    • AllGather 入力バッファのサイズは、モデルが小さいほど小さくなります。したがって、パラメータの数が少ないモデルでは、SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES に必要なサイズを小さくすることができます。

    • AllGather 入力バッファのサイズは、モデルがより多くの にシャーディングされるため、sharded_data_parallel_degree増加するにつれて小さくなりますGPUs。そのため、sharded_data_parallel_degree の値が大きいトレーニングジョブでは、SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES に必要なサイズを小さくすることができます。

  • SMDDP_AG_SORT_BUFFER_SIZE_BYTES のチューニング

    • ノード間通信から収集されるデータ量は、パラメータの数が少ないモデルほど少なくなります。したがって、パラメータの数が少ないこうしたモデルでは、SMDDP_AG_SORT_BUFFER_SIZE_BYTES に必要なサイズを小さくすることができます。

一部の集合体は を使用するようにフォールバックする可能性がありますNCCL。したがって、最適化されたSMDDP集合体からパフォーマンスの向上が得られない場合があります。追加のGPUメモリが使用可能な場合は、パフォーマンスの向上のメリットSMDDP_AG_SORT_BUFFER_SIZE_BYTESを得るために、 SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTESと の値を増やすことを検討できます。

次のコードは、 PyTorch 推定器のディストリビューションパラメータmpi_optionsで環境変数を に追加して環境変数を設定する方法を示しています。

import sagemaker from sagemaker.pytorch import PyTorch smp_options = { .... # All modelparallel configuration options go here } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8 # Required } # Use the following two lines to tune values of the environment variables for buffer mpioptions += " -x SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES=8192" mpioptions += " -x SMDDP_AG_SORT_BUFFER_SIZE_BYTES=8192" smd_mp_estimator = PyTorch( entry_point="your_training_script.py", # Specify your train script source_dir="location_to_your_script", role=sagemaker.get_execution_role(), instance_count=8, instance_type='ml.p4d.24xlarge', framework_version='1.13.1', py_version='py3', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="sharded-data-parallel-demo-with-tuning", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/')

シャーディングデータ並列処理による混合精度トレーニング

半精度の浮動小数点数とシャーディングデータ並列処理でGPUメモリをさらに節約するには、分散トレーニング設定に 1 つのパラメータを追加することで、16 ビット浮動小数点形式 (FP16) または動動小数点形式 (BF16) をアクティブ化できます。

注記

シャーディングデータ並列処理による混合精度トレーニングは、 SageMaker モデル並列処理ライブラリ v1.11.0 以降で利用できます。

シャーディングデータ並列処理によるFP16トレーニングの場合

シャーディングデータ並列処理を使用してFP16トレーニングを実行するには、 smp_optionsを設定ディクショナリ"fp16": True"に追加します。トレーニングスクリプトでは、smp.DistributedOptimizer モジュールを使用して静的損失スケーリングオプションと動的損失スケーリングオプションのいずれかを選択できます。詳細については、「FP16 モデル並列処理によるトレーニング」を参照してください。

smp_options = { "enabled": True, "parameters": { "ddp": True, "sharded_data_parallel_degree": 2, "fp16": True } }

シャーディングデータ並列処理によるBF16トレーニングの場合

のシャーディングデータ並列処理機能は、BF16データ型のトレーニング SageMaker をサポートします。BF16 データ型は 8 ビットを使用して浮動小数点数の指数を表し、FP16データ型は 5 ビットを使用します。指数に 8 ビットを保持すると、 は 32 ビット単精度浮動小数点 (FP32) 数の指数の同じ表現を維持できます。これにより、 FP32 と 間の変換BF16が簡単になり、特に大規模なモデルをトレーニングするときに、FP16トレーニングで頻繁に発生するオーバーフローやアンダーフローの問題を引き起こす可能性が大幅に低くなります。どちらのデータ型も合計で 16 ビットを使用しますが、この BF16形式の指数の表現範囲の増加は、精度の低下を犠牲にしています。大規模なモデルをトレーニングする場合、この精度の低下は、多くの場合、範囲とトレーニングの安定性にとって許容できるトレードオフと見なされます。

注記

現在、シャーディングデータ並列処理がアクティブ化されている場合にのみBF16トレーニングが機能します。

シャーディングデータ並列処理を使用してBF16トレーニングを実行するには、 smp_optionsを設定ディクショナリ"bf16": Trueに追加します。

smp_options = { "enabled": True, "parameters": { "ddp": True, "sharded_data_parallel_degree": 2, "bf16": True } }

テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理

シャーディングデータ並列処理を使用し、グローバルバッチサイズも小さくする必要がある場合は、テンソル並列処理とシャーディングデータ並列処理の使用を検討してください。非常に大きなコンピューティングクラスター (通常は 128 ノード以上) でシャーディングデータ並列処理を使用して大規模なモデルをトレーニングする場合、 あたりのバッチサイズが小さい場合でも、グローバルバッチサイズGPUが非常に大きくなります。これにより、収束の問題や計算パフォーマンスの低下の問題が発生する可能性があります。1 つのバッチがすでに大きく、それ以上減らすことができない場合、シャーディングデータ並列処理だけでは あたりのバッチサイズを減らすことができないGPUことがあります。このような場合、シャーディングデータ並列処理とテンソル並列処理を組み合わせて使用すると、グローバルバッチサイズを小さくすることができます。

最適なシャーディングデータ並列度とテンソル並列度の選択は、モデルの規模、インスタンスタイプ、およびモデルが収束するのに適したグローバルバッチサイズによって異なります。CUDA out-of-memory エラーを解決し、最高のパフォーマンスを実現するために、グローバルバッチサイズをコンピューティングクラスターに適合させるために、低テンソル並列度から始めることをお勧めします。テンソル並列処理とシャーディングデータ並列処理の組み合わせが、モデル並列処理GPUsのグループ化によるグローバルバッチサイズの調整にどのように役立つかについては、次の 2 つの例を参照してください。これにより、モデルレプリカの数が少なくなり、グローバルバッチサイズも小さくなります。

注記

この機能は SageMaker モデル並列処理ライブラリ v1.15 から使用でき、 PyTorch v1.13.1 をサポートしています。

注記

この機能は、ライブラリのテンソル並列処理機能によりサポート対象のモデルで使用できます。サポートされているモデルのリストについては、「Hugging Face Transformer モデルのサポート」を参照してください。また、トレーニングスクリプトを変更する際は、smp.model_creation 引数に tensor_parallelism=True を渡す必要があることにも注意してください。詳細については、SageMaker 「サンプル GitHub リポジトリtrain_gpt_simple.py」の「トレーニングスクリプト」を参照してください。

例 1

1536 GPUs (それぞれ 8 個のノードを持つ 192 GPUs個のノード) のクラスターでモデルをトレーニングするとします。シャーディングデータ並列処理の度合いを 32 (sharded_data_parallel_degree=32) に設定し、 あたりのバッチサイズを 1 GPUに設定します。各バッチのシーケンス長は 4096 トークンです。この場合、モデルレプリカは 1536 個、グローバルバッチサイズは 1536 個になり、各グローバルバッチには約 600 万個のトークンが含まれます。

(1536 GPUs) * (1 batch per GPU) = (1536 global batches) (1536 batches) * (4096 tokens per batch) = (6,291,456 tokens)

テンソル並列処理を追加すると、グローバルバッチサイズを小さくすることができます。設定例の 1 つは、テンソル並列度を 8 に設定し、 ごとのバッチサイズを 4 GPUに設定することです。これにより、192 個のテンソル並列グループまたは 192 個のモデルレプリカが形成され、各モデルレプリカは 8 個の に分散されますGPUs。バッチサイズ 4 は、反復ごと、およびテンソル並列グループごとのトレーニングデータの量です。つまり、各モデルレプリカは反復ごとに 4 つのバッチを消費します。この場合、グローバルバッチサイズは 768 になり、各グローバルバッチには約 300 万個のトークンが含まれます。したがって、シャーディングデータ並列処理のみを行った前のケースと比較して、グローバルバッチサイズは半分に削減されます。

(1536 GPUs) / (8 tensor parallel degree) = (192 tensor parallelism groups) (192 tensor parallelism groups) * (4 batches per tensor parallelism group) = (768 global batches) (768 batches) * (4096 tokens per batch) = (3,145,728 tokens)

例 2

シャーディングデータ並列処理とテンソル並列処理の両方が有効になると、ライブラリは最初にテンソル並列処理を適用し、このディメンション全体でモデルをシャードします。テンソル並列ランクごとに、データ並列処理は sharded_data_parallel_degree という形式で適用されます。

例えば、テンソル並列度を 4 (4 GPUs のグループを形成しGPUs)、シャーディングデータ並列度を 4 に設定し、レプリケーション度を 2 にするとします。この割り当てでは、テンソル並列度に基づいて、、(0,1,2,3)(4,5,6,7)、、、、(8,9,10,11)(12,13,14,15)、 の 8 つのGPUグループが作成されます(16,17,18,19)(20,21,22,23)(24,25,26,27)(28,29,30,31)。つまり、4 つは 1 つのテンソル並列グループGPUsを形成します。この場合、GPUsテンソル並列グループの 0 ランク目の縮小データ並列グループは になります(0,4,8,12,16,20,24,28)。縮小データ並列グループは、シャーディングデータ並列度 4 に基づいてシャーディングされ、データ並列処理用の 2 つのレプリケーショングループとなります。GPUs (0,4,8,12) は、0 番目のテンソル並列ランクのすべてのパラメータの完全なコピーをまとめて保持する 1 つのシャーディンググループを形成し、別のそのようなグループGPUs(16,20,24,28)を形成します。他のテンソル並列ランクにも同様のシャーディンググループとレプリケーショングループがあります。

図 1: テンソル並列処理グループ。

図 1: (ノード、シャーディングデータ並列度、テンソル並列度) = (4、4、4) のテンソル並列処理グループ。各長方形は 0~31 のインデックスGPUを持つ を表します。テンソル並列処理のGPUs形式は から TPG0ですTPG7。レプリケーショングループは ({TPG0, TPG4}, {TPG1, TPG5}, {TPG2, TPG6} および {, TPG7}) です。各レプリケーショングループペアは同じ色を共有しますがTPG3、塗り分けは異なります。

図 2: シャーディングデータ並列処理グループ。

図 2: (ノード、シャーディングデータ並列度、テンソル並列度) = (4、4、4) のシャーディングデータ並列処理グループ。各長方形は 0~31 のインデックスGPUを持つ を表します。から SDPG0へのシャードデータ並列処理グループというGPUs形式SDPG7。レプリケーショングループは ({SDPG0, SDPG4}, {SDPG1, SDPG5}, {SDPG2, SDPG6} および {, SDPG7}) です。各レプリケーショングループペアは同じ色を共有しますがSDPG3、塗り分けは異なります。

テンソル並列処理でシャーディングデータ並列処理を有効にする方法

テンソル並列処理でシャーディングデータ並列処理を使用するには、推定器クラスのオブジェクト SageMaker PyTorchを作成するdistributionときに、 の設定tensor_parallel_degreesharded_data_parallel_degreeと の両方を設定する必要があります。

また、prescaled_batch を有効化する必要があります。つまり、各テンソル並列グループは、独自のデータのバッチGPUを読み取る代わりに、選択したバッチサイズの結合バッチをまとめて読み取ります。実質的に、データセットを の数 GPUs (またはデータ並列サイズ、smp.dp_size()) と等しい部分に分割する代わりに、 の数を でGPUs割った数 tensor_parallel_degree (データ並列サイズの減少、) と等しい部分に分割しますsmp.rdp_size()。プリスケーリングされたバッチの詳細については、SageMaker Python SDKドキュメント「プリスケーリングされたバッチ」を参照してください。「サンプルリポジトリtrain_gpt_simple.py」のGPT「-2 のトレーニングスクリプトの例」も参照してください。 SageMaker GitHub

次のコードスニペットは、 の前述のシナリオに基づいて PyTorch 推定器オブジェクトを作成する例を示しています例 2

mpi_options = "-verbose --mca orte_base_help_aggregate 0 " smp_parameters = { "ddp": True, "fp16": True, "prescaled_batch": True, "sharded_data_parallel_degree": 4, "tensor_parallel_degree": 4 } pytorch_estimator = PyTorch( entry_point="your_training_script.py", role=role, instance_type="ml.p4d.24xlarge", volume_size=200, instance_count=4, sagemaker_session=sagemaker_session, py_version="py3", framework_version="1.13.1", distribution={ "smdistributed": { "modelparallel": { "enabled": True, "parameters": smp_parameters, } }, "mpi": { "enabled": True, "processes_per_host": 8, "custom_mpi_options": mpi_options, }, }, source_dir="source_directory_of_your_code", output_path=s3_output_location )

シャーディングデータ並列処理を使用する際のヒントと考慮事項

SageMaker モデル並列処理ライブラリのシャーディングデータ並列処理を使用する場合は、次の点を考慮してください。

  • シャーディングデータ並列処理はFP16トレーニングと互換性があります。FP16 トレーニングを実行するには、FP16 モデル並列処理によるトレーニング「」セクションを参照してください。

  • シャーディングデータ並列処理はテンソル並列処理と互換性があります。テンソル並列処理でシャーディングデータ並列処理を使用する場合に考慮する必要がある項目は次のとおりです。

    • テンソル並列処理とシャーディングデータ並列処理を一緒に使用すると、埋め込み層もテンソル並列グループ全体に自動分散されます。つまり、distribute_embedding パラメータは自動的に True に設定されます。テンセル並列処理の詳細については、「テンソル並列処理」を参照してください。

    • テンソル並列処理によるシャーディングデータ並列処理では、現在、分散トレーニング戦略のバックエンドとしてNCCL集合体を使用していることに注意してください。

    詳細については、「テンソル並列処理と NCCL Collectives によるシャーディングデータ並列処理」のセクションを参照してください。

  • 現在のところ、シャーディングデータ並列処理はパイプライン並列処理オプティマイザ状態シャーディングとは互換性がありません。シャーディングデータ並列処理を有効にするには、オプティマイザ状態のシャーディングをオフにし、パイプラインの並列度を 1 に設定します。

  • アクティベーションチェックポイント機能アクティベーションオフロード機能はシャーディングデータ並列処理と互換性があります。

  • 勾配累積でシャーディングデータ並列処理を使用するには、モデルを smdistributed.modelparallel.torch.DistributedModel モジュールでラップするときの累積ステップ数に backward_passes_per_step 引数を設定します。これにより、モデルレプリケーショングループ (シャーディンググループ) にまたがる勾配 AllReduce オペレーションが勾配累積の境界で確実に行われます。

  • シャーディングデータ並列処理でトレーニングされたモデルは、ライブラリのチェックポイント 、APIs、smp.save_checkpointおよび を使用してチェックポイントできますsmp.resume_from_checkpoint。詳細については、「分散 PyTorch モデルのチェックポイント ( SageMaker モデル並列処理ライブラリ v1.10.0 以降用)」を参照してください。

  • シャーディングデータ並列処理では、delayed_parameter_initialization 設定パラメータの挙動が変わります。これら 2 つの機能を同時にオンにすると、パラメータの初期化を遅らせるのではなく、モデル作成時にパラメータがすぐにシャーディング方式で初期化されるため、ランクごとに独自のパラメータのシャードが初期化されて保存されます。

  • シャーディングデータ並列処理を有効にすると、optimizer.step() の呼び出しの実行時にライブラリが内部で勾配クリッピングを実行します。などの勾配クリッピングAPIsに ユーティリティを使用する必要はありませんtorch.nn.utils.clip_grad_norm_()。勾配クリッピングのしきい値を調整するには、 シャーディングデータ並列処理をトレーニングジョブに適用する方法セクションに示すように、 SageMaker PyTorch 推定器を構築するときに、sdp_gradient_clippingディストリビューションパラメータ設定の パラメータを使用して設定できます。