Uso de la biblioteca de SMDDP en el script de entrenamiento de PyTorch Lightning - Amazon SageMaker AI

Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.

Uso de la biblioteca de SMDDP en el script de entrenamiento de PyTorch Lightning

Si quiere usar su script de entrenamiento PyTorch Lightning y ejecutar un trabajo de entrenamiento de paralelismo de datos distribuidos en SageMaker AI, puede ejecutar el trabajo de entrenamiento con cambios mínimos en el script de entrenamiento. Estos son los cambios necesarios que deben aplicarse: importar los módulos de PyTorch de la biblioteca de smdistributed.dataparallel, configurar las variables de entorno para que PyTorch Lightning acepte las variables de entorno de SageMaker AI predefinidas en el kit de herramientas de entrenamiento de SageMaker y activar la biblioteca de SMDDP 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 de paralelismo de datos de SageMaker AI v1.5.0 y versiones posteriores.

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

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

  2. Cree una instancia del 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 DDP de PyTorch: cree un objeto de la clase DDPStrategy con "smddp" para process_group_backend y "gpu" para accelerator y páselo 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: cree un objeto de la clase FSDPStrategy (con la política de encapsulado que prefiera) con "smddp" para process_group_backend y "gpu" para accelerator, y páselo a la clase 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
)

Cuando haya terminado de adaptar su script de entrenamiento, continúe con Inicialización de trabajos de entrenamiento distribuido con SMDDP mediante el SageMaker Python SDK.

nota

Al construir un estimador PyTorch de SageMaker AI y enviar una solicitud de trabajo de entrenamiento en Inicialización de trabajos de entrenamiento distribuido con SMDDP mediante el SageMaker Python SDK, debe proporcionar requirements.txt para instalar pytorch-lightning y lightning-bolts en el contenedor de entrenamiento PyTorch de SageMaker AI.

# requirements.txt
pytorch-lightning
lightning-bolts

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