Use a biblioteca SMDDP em seu script de treinamento do PyTorch Lightning - SageMaker IA da Amazon

As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

Use a biblioteca SMDDP em seu script de treinamento do PyTorch Lightning

Se quiser trazer seu script de treinamento do PyTorch Lightning e executar uma tarefa de treinamento paralelo de dados distribuídos no SageMaker AI, você pode executar a tarefa de treinamento com o mínimo de alterações em seu script de treinamento. As alterações necessárias incluem o seguinte: importar os módulos do PyTorch da biblioteca smdistributed.dataparallel, configurar as variáveis de ambiente do PyTorch Lightning para aceitar as variáveis de ambiente do SageMaker AI que estão predefinidas no kit de ferramentas de treinamento do SageMaker e ativar a biblioteca de SMDDP definindo o backend do grupo de processos como "smddp". Para saber mais, siga as instruções a seguir que detalham as etapas com exemplos de código.

nota

É possível usar o PyTorch Lightning na biblioteca de paralelismo de dados do SageMaker AI v1.5.0 e versões posteriores.

  1. Importe a biblioteca pytorch_lightning e os módulos smdistributed.dataparallel.torch.

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

  2. Instancie o 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: Crie um objeto da classe DDPStrategy com "smddp" para process_group_backend e "gpu" para accelerator e passe-o para a 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
)

    Para PyTorch FSDP: Crie um objeto da classe FSDPStrategy (com a política de empacotamento de sua escolha) com "smddp" para process_group_backend e "gpu" para accelerator passe isso para a 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
)

Depois de concluir a adaptação do seu script de treinamento, prossiga para Lançamento de trabalhos de treinamento distribuídos com SMDDP com o uso do SageMaker Python SDK.

nota

Ao criar um estimador do PyTorch do SageMaker AI e enviar uma solicitação de tarefa de treinamento na Lançamento de trabalhos de treinamento distribuídos com SMDDP com o uso do SageMaker Python SDK, você precisa fornecer requirements.txt para instalar pytorch-lightning e lightning-bolts no contêiner de treinamento do PyTorch do SageMaker AI.

# requirements.txt
pytorch-lightning
lightning-bolts

Para ter mais informações sobre como especificar o diretório de origem para colocar o arquivo requirements.txt com seu script de treinamento e como enviar uma tarefa, consulte Using third-party libraries na documentação do Amazon SageMaker AI Python SDK.