Usa la libreria SMDDP nello script di addestramento Lightning PyTorch

Se desideri utilizzare lo script di formazione PyTorchLightning ed eseguire un processo di formazione parallelo con dati distribuiti SageMaker, puoi eseguire il processo di formazione apportando modifiche minime allo script di formazione. Le modifiche necessarie includono quanto segue: importa i PyTorch moduli della smdistributed.dataparallel libreria, imposta le variabili di ambiente affinché PyTorch Lightning accetti le variabili di SageMaker ambiente preimpostate dal toolkit di SageMaker formazione e attiva la libreria SMDDP impostando il backend del gruppo di processi su. "smddp" Per ulteriori informazioni, segui le seguenti istruzioni che descrivono le fasi con esempi di codice.

Nota

Il supporto PyTorch Lightning è disponibile nella libreria SageMaker data parallel v1.5.0 e versioni successive.

PyTorch Lightning == v2.1.0 e == 2.0.1 PyTorch

1. Importa la libreria pytorch_lightning e i moduli smdistributed.dataparallel.torch.

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

2. Crea un'LightningEnvironmentistanza di.

   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. Per PyTorch DDP: crea un oggetto della classe DDPStrategy con "smddp" for process_group_backend e "gpu" for accelerator e passalo alla 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
   )

   Per PyTorch FSDP: crea un oggetto della classe FSDPStrategy (con la politica di wrapping preferita) con "smddp" for process_group_backend e for e "gpu" passalo alla classe Trainer. accelerator

   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
   )

Dopo aver completato l'adattamento dello script di addestramento, procedi con Passaggio 2: Avvia un processo di formazione distribuito utilizzando SageMaker Python SDK.

Nota

Quando crei uno SageMaker PyTorch estimatore e invii una richiesta di lavoro di formazione inPassaggio 2: Avvia un processo di formazione distribuito utilizzando SageMaker Python SDK, devi provvedere all'installazione e inserirla nel contenitore di formazione. requirements.txt pytorch-lightning lightning-bolts SageMaker PyTorch

# requirements.txt
pytorch-lightning
lightning-bolts

Per ulteriori informazioni su come specificare la directory di origine in cui inserire il requirements.txt file insieme allo script di formazione e all'invio di un lavoro, consulta la sezione Utilizzo di librerie di terze parti nella documentazione di Amazon SageMakerPython SDK.