ハイブリッドシャードデータ並列処理

シャードデータ並列処理は、モデルの状態 (モデルパラメータ、勾配、オプティマイザの状態) をデバイス間で分割する、メモリを節約する分散トレーニング手法です。これにより、解放されたGPUメモリを使用して、より大きなモデルを適合させたり、バッチサイズを増やしたりすることができます。SMP ライブラリは、 PyTorch フルシャードデータ並列 (FSDP） PyTorch FSDP でシャードデータ並列処理を実行する機能を提供します。デフォルトでは、使用されている セット全体でシャードGPUsが実行されます。SMP v2 では、ライブラリは PyTorch 、ハイブリッドシャーディング (HYBRID_SHARD) を拡張することで、このシャーディングデータ並列処理を に加えて PyTorch FSDP提供します。これは、、FULL_SHARD、SHARD_GRAD_OP、HYBRID_SHARD、 が提供するシャーディング戦略 PyTorch FSDPの 1 つです_HYBRID_SHARD_ZERO2。この方法でハイブリッドシャーディングを拡張すると、 の PyTorchブログ「 での巨大なモデルトレーニングのほぼ線形スケーリング AWS」で説明されているように、 を実装 scale-aware-shardingするのに役立ちますFSDP。

SMP ライブラリを使用すると、設定可能な任意の数の _HYBRID_SHARD_ZERO2で HYBRID_SHARDおよび を簡単に使用できGPUs、1 つのノード (HYBRID_SHARD) またはすべての GPUs () でシャーディングをサポートするネイティブ PyTorch FSDPを拡張できますFULL_SHARD。 PyTorch FSDP 呼び出しはそのままにしておくことができ、次のコード例に示すように、SMP設定にhybrid_shard_degree引数を追加するだけで済みます。 PyTorch モデルのFSDPラッパー内のsharding_strategy PyTorch引数の値を変更する必要はありません。を値ShardingStrategy.HYBRID_SHARDとして渡すことができます。または、 hybrid_shard_degree パラメータに 2 以上の値を指定ShardingStrategy.HYBRID_SHARDした場合、SMPライブラリはスクリプトの戦略を上書きし、 に設定します。

次のコードスニペットは、 で導入された 2 ステッププロセスに従って、SMP初期化モジュールtorch.sagemaker.init()をトレーニングスクリプトに追加し、トレーニングジョブランチャー用の JSON 形式でSMP設定ディクショナリを設定する方法を示しています SageMaker モデル並列処理ライブラリ v2 を使用する。 PyTorch モデルやPyTorch FSDP設定を変更する必要はありません。hybrid_shard_degree パラメータの詳細については、「SMP v2 コア機能設定パラメータ」をご参照ください。

SMP 設定ディクショナリ

{ "hybrid_shard_degree": 16 }

トレーニングスクリプト内

import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
    model,
    ...
)

# Optimizer needs to be created after FSDP wrapper
optimizer = ...