Utilice la biblioteca SMDDP en su PyTorch script de entrenamiento de 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.

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

Si desea utilizar su guion de entrenamiento de PyTorchLightning y ejecutar un trabajo de formación en paralelo de datos distribuidos en SageMaker IA, puede ejecutar el trabajo de formación con cambios mínimos en su guion de formación. 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 entorno de SageMaker IA predefinidas en el kit de herramientas de SageMaker entrenamiento y activar la biblioteca 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 paralela de datos de SageMaker IA 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. 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 el PyTorch DDP: crea un objeto de la DDPStrategyclase con el para y el "smddp""gpu" paraaccelerator, process_group_backend 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 el PyTorch FSDP: crea un objeto de la FSDPStrategyclase (con la política de empaquetado que prefieras) con un para y un "smddp""gpu" para process_group_backend y accelerator 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 con SMDDP mediante el SDK de Python SageMaker .

nota

Cuando crees un PyTorch estimador de SageMaker IA y envíes una solicitud de trabajo de formaciónLanzamiento de trabajos de formación distribuidos con SMDDP mediante el SDK de Python SageMaker , tendrás que proporcionar la instalación pytorch-lightning y el requirements.txt SageMaker contenedor de formación lightning-bolts en IA. PyTorch 

# 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 entrenamiento y el envío de un trabajo, consulte Uso de bibliotecas de terceros en la documentación del SDK de Python de Amazon SageMaker AI.