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

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

注記

PyTorch Lightning サポートは、 SageMaker データ並列ライブラリ 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. For PyTorch DDP — for process_group_backendと "smddp" for "gpu" を使用してDDPStrategyクラスのオブジェクトを作成しaccelerator、それをトレーナークラスに渡します。

   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
   )

   For PyTorch FSDP — for process_group_backendと "smddp" for を使用してFSDPStrategyクラスのオブジェクト (選択したラッピングポリシーを使用) "gpu" を作成しaccelerator、それをトレーナークラスに渡します。

   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 SMDDPを使用して で分散トレーニングジョブを起動する SDK に進みます。

注記

SageMaker PyTorch 推定器を作成し、 でトレーニングジョブリクエストを送信する場合は SageMaker Python SMDDPを使用して で分散トレーニングジョブを起動する SDK、 を PyTorch トレーニングコンテナpytorch-lightninglightning-boltsに SageMakerインストールrequirements.txtして を指定する必要があります。

# requirements.txt
pytorch-lightning
lightning-bolts

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