Utilice la SMDDP biblioteca en su script de entrenamiento de PyTorch Lightning

Si desea incorporar su script de entrenamiento PyTorchLightning y ejecutar un trabajo de entrenamiento paralelo de datos distribuidos SageMaker, puede ejecutar el trabajo de entrenamiento con cambios mínimos en su guión de entrenamiento. Los cambios necesarios incluyen los siguientes: importar los PyTorch módulos de la smdistributed.dataparallel biblioteca, configurar las variables de entorno para que PyTorch Lightning acepte las variables de SageMaker entorno predefinidas en el kit de herramientas de SageMaker entrenamiento y activar la SMDDP biblioteca configurando el backend del grupo de procesos en. "smddp" Para obtener más información, siga las siguientes instrucciones que desglosan los pasos con ejemplos de código.

nota

La compatibilidad con PyTorch Lightning está disponible en la biblioteca paralela de SageMaker datos v1.5.0 y versiones posteriores.

PyTorch Lightning == v2.1.0 y == 2.0.1 PyTorch

1. Importe la biblioteca pytorch_lightning y los módulos smdistributed.dataparallel.torch.

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

2. Instancie el. 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. Para PyTorch DDP: crea un objeto de la DDPStrategyclase con "smddp" para process_group_backend y "gpu" paraaccelerator, y pásalo a la clase 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
   )

   Para PyTorch FSDP: crea un objeto de la FSDPStrategyclase (con la política de empaquetado que prefieras) con un "smddp" por process_group_backend y un "gpu" para accelerator y pásalo a la clase de entrenador.

   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
   )

Cuando haya terminado de adaptar su script de entrenamiento, continúe con Lanzamiento de trabajos de formación distribuidos SMDDP con SageMaker Python SDK.

nota

Cuando se crea un SageMaker PyTorch estimador y se envía una solicitud de trabajo de formaciónLanzamiento de trabajos de formación distribuidos SMDDP con SageMaker Python SDK, se debe requirements.txt proporcionar la instalación pytorch-lightning y el contenedor lightning-bolts de SageMaker PyTorch formación.

# requirements.txt
pytorch-lightning
lightning-bolts

Para obtener más información sobre cómo especificar el directorio de origen para colocar el requirements.txt archivo junto con el guion de formación y el envío de un trabajo, consulte Uso de bibliotecas de terceros en la SDKdocumentación de Amazon SageMaker Python.