Use os estimadores da estrutura de PyTorch no SageMaker Python SDK - 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 os estimadores da estrutura de PyTorch no SageMaker Python SDK

Você pode iniciar o treinamento distribuído adicionando o argumento distribution ao estimador PyTorch ou TensorFlow do framework do SageMaker AI. Para obter mais detalhes, escolha um dos frameworks compatíveis com a biblioteca de paralelismo de dados distribuídos do SageMaker AI (SMDDP) das seguintes seleções:

PyTorch

As seguintes opções de inicializador estão disponíveis para iniciar o treinamento distribuído do PyTorch:

  • pytorchddp: essa opção executa mpirun e configura as variáveis de ambiente necessárias para executar o treinamento distribuído do PyTorch no SageMaker AI. Para usar essa opção, transfira o seguinte dicionário para o parâmetro distribution:

    { "pytorchddp": { "enabled": True } }

  • torch_distributed: essa opção executa torchrun e configura as variáveis de ambiente necessárias para executar o treinamento distribuído do PyTorch no SageMaker AI. Para usar essa opção, transfira o seguinte dicionário para o parâmetro distribution:

    { "torch_distributed": { "enabled": True } }

  • smdistributed: essa opção também executa mpirun, mas com smddprun, que configura as variáveis de ambiente necessárias para executar o treinamento distribuído do PyTorch no SageMaker AI.

    { "smdistributed": { "dataparallel": { "enabled": True } } }

Se você optar por substituir o NCCL AllGather pelo SMDDP AllGather, poderá usar todas as três opções. Escolha uma opção que seja mais adequada ao caso de uso.

Se você optar por substituir o NCCL AllReduce pelo SMDDP AllReduce, deverá escolher uma das opções encontradas em mpirun: smdistributed ou pytorchddp. Você também pode adicionar outras opções de MPI da seguinte forma:

{ 
    "pytorchddp": {
        "enabled": True, 
        "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
    }
}
{ 
    "smdistributed": { 
        "dataparallel": {
            "enabled": True, 
            "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
        }
    }
}

O exemplo de código a seguir mostra a estrutura básica de um estimador PyTorch com opções de treinamento distribuídas.

from sagemaker.pytorch import PyTorch

pt_estimator = PyTorch(
    base_job_name="training_job_name_prefix",
    source_dir="subdirectory-to-your-code",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    py_version="py310",
    framework_version="2.0.1",

    # For running a multi-node distributed training job, specify a value greater than 1
    # Example: 2,3,4,..8
    instance_count=2,

    # Instance types supported by the SageMaker AI data parallel library: 
    # ml.p4d.24xlarge, ml.p4de.24xlarge
    instance_type="ml.p4d.24xlarge",

    # Activate distributed training with SMDDP
    distribution={ "pytorchddp": { "enabled": True } }  # mpirun, activates SMDDP AllReduce OR AllGather
    # distribution={ "torch_distributed": { "enabled": True } }  # torchrun, activates SMDDP AllGather
    # distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }  # mpirun, activates SMDDP AllReduce OR AllGather
)

pt_estimator.fit("s3://bucket/path/to/training/data")
nota

O PyTorch Lightning e as respectivas bibliotecas de utilitários, como o Lightning Bolts, não estão pré-instalados nos DLCs do PyTorch do SageMaker AI. Crie o arquivo requirements.txt a seguir e salve no diretório de origem em que você salva o script de treinamento.

# requirements.txt
pytorch-lightning
lightning-bolts

Por exemplo, o diretório de árvore estruturada deve ser semelhante ao seguinte:

├── pytorch_training_launcher_jupyter_notebook.ipynb
└── sub-folder-for-your-code
    ├──  adapted-training-script.py
    └──  requirements.txt

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.

Considerações para ativar as operações coletivas do SMDDP e usar as opções corretas de inicializador de treinamento distribuído

  • SMDDP AllReduce e SMDDP AllGather não são mutuamente compatíveis no momento.

  • O SMDDP AllReduce é ativado por padrão ao usar smdistributed ou pytorchddp, que são iniciadores encontrados no mpirun, e é usado o NCCL AllGather.

  • O SMDDP AllGather é ativado por padrão ao usar o iniciador torch_distributed e AllReduce volta para o NCCL.

  • O SMDDP AllGather pode ser ativado também com o uso dos iniciadores encontrados em mpirun com uma variável de ambiente adicional definida da seguinte forma:

    export SMDATAPARALLEL_OPTIMIZE_SDP=true
TensorFlow
Importante

A biblioteca SMDDP interrompeu o compatibilidade com o TensorFlow e não está mais disponível em DLCs para TensorFlow posteriores à v2.11.0. Para encontrar DLCs anteriores do TensorFlow com a biblioteca SMDDP instalada, consulte TensorFlow (obsoleto).

from sagemaker.tensorflow import TensorFlow

tf_estimator = TensorFlow(
    base_job_name = "training_job_name_prefix",
    entry_point="adapted-training-script.py",
    role="SageMakerRole",
    framework_version="2.11.0",
    py_version="py38",

    # For running a multi-node distributed training job, specify a value greater than 1
    # Example: 2,3,4,..8
    instance_count=2,

    # Instance types supported by the SageMaker AI data parallel library: 
    # ml.p4d.24xlarge, ml.p3dn.24xlarge, and ml.p3.16xlarge
    instance_type="ml.p3.16xlarge",

    # Training using the SageMaker AI data parallel distributed training strategy
    distribution={ "smdistributed": { "dataparallel": { "enabled": True } } }
)

tf_estimator.fit("s3://bucket/path/to/training/data")