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 você quiser trazer seu script de treinamento do PyTorchLightning e executar um trabalho de treinamento paralelo de dados distribuídos em SageMaker IA, você pode executar o trabalho de treinamento com o mínimo de alterações em seu script de treinamento. As mudanças necessárias incluem o seguinte: importar os PyTorch módulos da smdistributed.dataparallel biblioteca, configurar as variáveis de ambiente para que o PyTorch Lightning aceite as variáveis de ambiente de SageMaker IA predefinidas pelo kit de ferramentas de SageMaker treinamento e ative a biblioteca SMDDP configurando o back-end 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

O suporte ao PyTorch Lightning está disponível na biblioteca paralela de dados 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 DDPStrategyclasse com "smddp" for process_group_backend e for e accelerator passe-o "gpu" 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 FSDPStrategyclasse (com a política de empacotamento de escolha) com "smddp" for process_group_backend e for e passe isso "gpu" para accelerator 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 usando o Python SDK SageMaker .

nota

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

# requirements.txt
pytorch-lightning
lightning-bolts

Para obter mais informações sobre como especificar o diretório de origem para colocar o requirements.txt arquivo junto com seu script de treinamento e o envio de um trabalho, consulte Uso de bibliotecas de terceiros na documentação do Amazon SageMaker AI Python SDK.