Utilisation de la bibliothèque SMDDP dans un 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.

Utilisation de la bibliothèque SMDDP dans un 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 de parallélisme distribué des données dans SageMaker AI, vous pouvez exécuter la tâche d’entraînement en apportant un minimum de modifications à votre script d’entraînement. Les modifications nécessaires incluent : importation des modules PyTorch de la bibliothèque smdistributed.dataparallel, configuration des variables d’environnement pour que PyTorch Lightning accepte les variables d’environnement SageMaker prédéfinies par la boîte à outils d’entraînement SageMaker, et activation de la bibliothèque SMDDP en définissant le système dorsal 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 SageMaker AI de parallélisme des données version v1.5.0 et ultérieure.

  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 l’environnement Lightning (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" pour process_group_backend et "gpu" pour accelerator, 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" pour process_group_backend et "gpu" pour accelerator, 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 à Lancement de tâches d’entraînement distribué avec SMDDP à l’aide du kit SageMaker Python SDK.

Note

Lorsque vous construisez un estimateur SageMaker PyTorch et soumettez une demande de tâche d’entraînement dans Lancement de tâches d’entraînement distribué avec SMDDP à l’aide du kit SageMaker Python SDK, vous devez fournir requirements.txt pour installer pytorch-lightning et lightning-bolts dans le conteneur d’entraînement SageMaker AI PyTorch.

# requirements.txt
pytorch-lightning
lightning-bolts

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