Utiliser la bibliothèque SMDDP dans votre script d'entraînement PyTorch Lightning

Si vous souhaitez intégrer votre script d'entraînement PyTorchLightning et exécuter une tâche de formation parallèle aux données distribuées SageMaker, vous pouvez exécuter la tâche de formation en apportant des modifications minimales à votre script de formation. Les modifications nécessaires sont les suivantes : importation des PyTorch modules de la smdistributed.dataparallel bibliothèque, configuration des variables d'environnement pour que PyTorch Lightning accepte les variables d' SageMaker environnement prédéfinies par le kit de SageMaker formation, et activation de la bibliothèque SMDDP en définissant le backend du groupe de processus sur. "smddp" Pour en savoir plus, suivez les instructions ci-dessous qui décomposent les étapes avec des exemples de code.

Note

Le support PyTorch Lightning est disponible dans la bibliothèque SageMaker Data Parallel v1.5.0 et versions ultérieures.

PyTorch Lightning == v2.1.0 et PyTorch == 2.0.1

1. Importez la bibliothèque pytorch_lightning et les modules smdistributed.dataparallel.torch.

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

2. Instanciez le. 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. Pour PyTorch DDP : créez un objet de la classe DDPStrategy avec "smddp" for process_group_backend et "gpu" foraccelerator, et transmettez-le à la classe 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
   )

   Pour PyTorch FSDP : créez un objet de la classe FSDPStrategy (avec la politique d'encapsulation de votre choix) avec "smddp" for process_group_backend et "gpu" foraccelerator, et transmettez-le à la classe 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
   )

Une fois que vous avez terminé d'adapter votre scénario d'entraînement, passez à Étape 2 : Lancer une tâche de formation distribuée à l'aide du SDK SageMaker Python.

Note

Lorsque vous créez un SageMaker PyTorch estimateur et que vous soumettez une demande de formationÉtape 2 : Lancer une tâche de formation distribuée à l'aide du SDK SageMaker Python, vous devez fournir l'installation pytorch-lightning et requirements.txt l'inclure lightning-bolts dans le conteneur de SageMaker PyTorch formation.

# requirements.txt
pytorch-lightning
lightning-bolts

Pour plus d'informations sur la spécification du répertoire source dans lequel placer le requirements.txt fichier avec votre script d'entraînement et la soumission d'une tâche, consultez la section Utilisation de bibliothèques tierces dans la documentation du SDK Amazon SageMaker Python.