PyTorch Lightning トレーニングスクリプトで SMDDP ライブラリを使用する

PyTorch Lightning トレーニングスクリプトを持ち込んで SageMaker AI で分散データ並列トレーニングジョブを実行する場合は、トレーニングスクリプトの変更を最小限に抑えてトレーニングジョブを実行できます。必要な変更には、smdistributed.dataparallelライブラリの PyTorch モジュールをインポートし、SageMaker トレーニングツールキットによって事前設定された SageMaker SageMaker AI 環境変数を受け入れるように PyTorch Lightning の環境変数を設定し、プロセスグループのバックエンドを に設定して SMDDP ライブラリをアクティブ化することが含まれます"smddp"。詳細については、コード例を使って各ステップごとに説明します。

注記

PyTorch Lightning のサポートは、SageMaker AI データ並列ライブラリ v1.5.0 以降で利用できます。

PyTorch Lightning == v2.1.0 および PyTorch == 2.0.1

1. pytorch_lightning ライブラリと smdistributed.dataparallel.torch モジュールをインポートします。

   import lightning as pl
   import smdistributed.dataparallel.torch.torch_smddp

2. LightningEnvironment をインスタンス化します。

   from lightning.fabric.plugins.environments.lightning import LightningEnvironment
   
   env = LightningEnvironment()
   env.world_size = lambda: int(os.environ["WORLD_SIZE"])
   env.global_rank = lambda: int(os.environ["RANK"])

3. PyTorch DDP の場合 – DDPStrategy クラスのオブジェクトを作成し、process_group_backend には "smddp"、accelerator には "gpu" を指定します。このオブジェクトを Trainer クラスに渡します。

   import lightning as pl
   from lightning.pytorch.strategies import DDPStrategy
   
   ddp = DDPStrategy(
       cluster_environment=env, 
       process_group_backend="smddp", 
       accelerator="gpu"
   )
   
   trainer = pl.Trainer(
       max_epochs=200, 
       strategy=ddp, 
       devices=num_gpus, 
       num_nodes=num_nodes
   )

   PyTorch FSDP の場合 – FSDPStrategy クラスのオブジェクトを (選択したラッピングポリシーで) 作成し、process_group_backend には "smddp"、accelerator には "gpu" を指定します。このオブジェクトを Trainer クラスに渡します。

   import lightning as pl
   from lightning.pytorch.strategies import FSDPStrategy
   
   from functools import partial
   from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
   
   policy = partial(
       size_based_auto_wrap_policy, 
       min_num_params=10000
   )
   
   fsdp = FSDPStrategy(
       auto_wrap_policy=policy,
       process_group_backend="smddp", 
       cluster_environment=env
   )
   
   trainer = pl.Trainer(
       max_epochs=200, 
       strategy=fsdp, 
       devices=num_gpus, 
       num_nodes=num_nodes
   )

トレーニングスクリプトの調整が完了したら、SageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する に進みます。

注記

SageMaker AI PyTorch 推定器を構築し、 でトレーニングジョブリクエストを送信するときはSageMaker Python SDK を使用して SMDDP で分散トレーニングジョブを開始する、SageMaker AI PyTorch トレーニングコンテナlightning-boltsに pytorch-lightningと をインストールrequirements.txtするために を指定する必要があります。

# requirements.txt
pytorch-lightning
lightning-bolts

トレーニングスクリプトとジョブ送信とともにrequirements.txtファイルを配置するソースディレクトリを指定する方法の詳細については、Amazon SageMaker AI Python SDK ドキュメントの「サードパーティーライブラリの使用」を参照してください。