View a markdown version of this page

Utilisez la bibliothèque SMDDP dans votre script d'entraînement PyTorch Lightning - Amazon SageMaker AI

Les traductions sont fournies par des outils de traduction automatique. En cas de conflit entre le contenu d'une traduction et celui de la version originale en anglais, la version anglaise prévaudra.

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

Si vous souhaitez apporter votre script d'entraînement PyTorch Lightning et exécuter une tâche d'entraînement parallèle aux données distribuées dans SageMaker AI, vous pouvez exécuter la tâche de formation en modifiant le moins possible votre script d'entraînement. Les modifications nécessaires sont notamment les suivantes : importez les PyTorch modules de la smdistributed.dataparallel bibliothèque, configurez les variables d'environnement pour que PyTorch Lightning accepte les variables d'environnement SageMaker AI prédéfinies par la boîte à outils de SageMaker formation, et activez la bibliothèque SMDDP en réglant 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

La prise en charge de PyTorch Lightning est disponible dans la bibliothèque parallèle de données SageMaker AI v1.5.0 et versions ultérieures.

  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. https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html

    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 à Lancement de tâches de formation distribuées avec SMDDP à l'aide du SDK Python SageMaker.

Note

Lorsque vous créez une SageMaker IA PyTorch ModelTrainer et que vous soumettez une demande de formation dansLancement de tâches de formation distribuées avec SMDDP à l'aide du SDK Python SageMaker, vous devez fournir requirements.txt à installer pytorch-lightning et lightning-bolts dans le conteneur de PyTorch formation SageMaker AI.

# 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 de formation et la soumission d'une tâche, consultez la section Utilisation de bibliothèques tierces dans la documentation du SDK Amazon SageMaker AI Python.