Utilizzare la libreria SMDDP in uno script di addestramento PyTorch Lighting - Amazon SageMaker AI

Le traduzioni sono generate tramite traduzione automatica. In caso di conflitto tra il contenuto di una traduzione e la versione originale in Inglese, quest'ultima prevarrà.

Utilizzare la libreria SMDDP in uno script di addestramento PyTorch Lighting

Se desideri utilizzare il tuo script di addestramento PyTorch Lightning ed eseguire un job di addestramento Distributed Data Parallelism in SageMaker AI, puoi eseguire il job di addestramento con modifiche minime allo script. Le modifiche necessarie includono quanto segue: importare i moduli PyTorch della libreria smdistributed.dataparallel, impostare le variabili di ambiente per PyTorch Lightning in modo che accetti le variabili di ambiente SageMaker AI preimpostate dal Kit di strumenti di SageMaker Training e attivare 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 per PyTorch Lightning è disponibile nella libreria SageMaker AI Distributed Data Parallelism v1.5.0 e versioni successive.

  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’istanza di 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. Per PyTorch DDP - Crea un oggetto della classe DDPStrategy con "smddp" per process_group_backend e "gpu" per accelerator, quindi 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 una policy di wrapping qualsiasi) con "smddp" per process_group_backend e "gpu" per accelerator, quindi passalo alla 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
)

Dopo aver completato l’adattamento dello script di addestramento, procedi con Avvio di job di addestramento distribuito con SMDDP mediante SageMaker Python SDK.

Nota

Quando crei uno strumento di stima PyTorch di SageMaker AI e invii una richiesta di job di addestramento in Avvio di job di addestramento distribuito con SMDDP mediante SageMaker Python SDK, devi fornire requirements.txt per installare pytorch-lightning e lightning-bolts all’interno del container di addestramento PyTorch di SageMaker AI.

# requirements.txt
pytorch-lightning
lightning-bolts

Per ulteriori informazioni su come specificare la directory di origine in cui inserire il file requirements.txt insieme allo script di addestramento e all’invio di un processo, consulta Using third-party libraries nella documentazione di Amazon SageMaker AI Python SDK.