Guia de treinamento de GPU distribuída (SDK v2)

APLICA-SE A: SDK do Python azure-ai-ml v2 (atual)

Saiba mais sobre como usar o código de treinamento distribuído entre GPUs no Azure Machine Learning. Este artigo ajuda você a executar seu código de treinamento distribuído existente e oferece dicas e exemplos para você seguir relativos a cada estrutura:

• PyTorch
• TensorFlow
• Acelere o treinamento de GPU com o InfiniBand

Pré-requisitos

Leia esses conceitos básicos de treinamento distribuído entre GPUs, como paralelismo de dados, paralelismo de dados distribuídos e paralelismo de modelos.

PyTorch

O Azure Machine Learning dá suporte à execução de trabalhos distribuídos usando as funcionalidades nativas de treinamento distribuído do PyTorch (torch.distributed).

Inicialização do grupo de processos

O backbone de um treinamento distribuído é baseado em um grupo de processos que se conhecem e podem se comunicar entre si por um back-end. No PyTorch, o grupo de processos é criado ao chamar torch.distributed.init_process_group em todos os processos distribuídos para formar coletivamente um grupo do processo.

torch.distributed.init_process_group(backend='nccl', init_method='env://', ...)

Os back-ends de comunicação mais comuns usados são mpi, nccl e gloo. Para o treinamento baseado em GPU, a nccl é recomendada para o melhor desempenho e deve ser usada sempre que possível.

init_method informa como cada processo pode descobrir um ao outro, como eles inicializam e verificar o grupo de processos usando o back-end de comunicação. Por padrão, se o init_method não for especificado, o PyTorch usará o método de inicialização de variável de ambiente (env://). init_method é o método de inicialização recomendado a ser usado no código de treinamento para executar o PyTorch distribuído no Azure Machine Learning. O PyTorch irá procurar as variáveis de ambiente a seguir para a inicialização:

• MASTER_ADDR: endereço IP do computador que hospeda o processo com a classificação 0
• MASTER_PORT: uma porta livre no computador que hospeda o processo com a classificação 0
• WORLD_SIZE: o número total de processos. Deve ser igual ao número total de dispositivos (GPUs) usados para o treinamento distribuído
• RANK: a classificação (global) do processo atual. Os valores possíveis são 0 a (tamanho do mundo: 1)

Para obter mais informações sobre a inicialização do grupo de processos, confira a documentação do PyTorch.

Muitos aplicativos também precisam das seguintes variáveis de ambiente:

• LOCAL_RANK: a classificação local (relativa) do processo dentro do nó. Os valores possíveis são 0 a (n º de processos no nó - 1). Essas informações são úteis porque muitas operações, como a preparação de dados, só devem ser executadas uma vez por nó, geralmente no local_rank = 0.
• NODE_RANK: a classificação do nó para treinamento em vários nós. Os valores possíveis são 0 a (n º total de nós - 1).

Você não precisa usar um utilitário inicializador como o torch.distributed.launch. Para executar um trabalho distribuído do PyTorch:

1. Especifique os argumentos e o script de treinamento.
2. Crie um command e especifique o tipo como PyTorch e o process_count_per_instance no parâmetro distribution. O process_count_per_instance corresponde ao número total de processos que você deseja executar para o trabalho. process_count_per_instance normalmente deve ser igual a # of GPUs per node. Se process_count_per_instance não for especificado, por padrão, o Azure Machine Learning iniciará um processo por nó.

O Azure Machine Learning define as variáveis de ambiente MASTER_ADDR, MASTER_PORT, WORLD_SIZE e NODE_RANK em cada nó, o nível do processo RANK e as variáveis de ambiente de LOCAL_RANK.

from azure.ai.ml import command
from azure.ai.ml.entities import Data
from azure.ai.ml import Input
from azure.ai.ml import Output
from azure.ai.ml.constants import AssetTypes

# === Note on path ===
# can be can be a local path or a cloud path. AzureML supports https://`, `abfss://`, `wasbs://` and `azureml://` URIs.
# Local paths are automatically uploaded to the default datastore in the cloud.
# More details on supported paths: https://docs.microsoft.com/azure/machine-learning/how-to-read-write-data-v2#supported-paths

inputs = {
    "cifar": Input(
        type=AssetTypes.URI_FOLDER, path=returned_job.outputs.cifar.path
    ),  # path="azureml:azureml_stoic_cartoon_wgb3lgvgky_output_data_cifar:1"), #path="azureml://datastores/workspaceblobstore/paths/azureml/stoic_cartoon_wgb3lgvgky/cifar/"),
    "epoch": 10,
    "batchsize": 64,
    "workers": 2,
    "lr": 0.01,
    "momen": 0.9,
    "prtfreq": 200,
    "output": "./outputs",
}

from azure.ai.ml.entities import ResourceConfiguration

job = command(
    code="./src",  # local path where the code is stored
    command="python train.py --data-dir ${{inputs.cifar}} --epochs ${{inputs.epoch}} --batch-size ${{inputs.batchsize}} --workers ${{inputs.workers}} --learning-rate ${{inputs.lr}} --momentum ${{inputs.momen}} --print-freq ${{inputs.prtfreq}} --model-dir ${{inputs.output}}",
    inputs=inputs,
    environment="azureml:AzureML-acpt-pytorch-2.2-cuda12.1@latest",
    instance_count=2,  # In this, only 2 node cluster was created.
    distribution={
        "type": "PyTorch",
        # set process count to the number of gpus per node
        # NC6s_v3 has only 1 GPU
        "process_count_per_instance": 1,
    },
)
job.resources = ResourceConfiguration(
    instance_type="Standard_NC6s_v3", instance_count=2
)  # Serverless compute resources

Exemplo de Pytorch

• Para que o notebook completo execute o exemplo do Pytorch, confira azureml-examples: Treinamento distribuído com PyTorch no CIFAR-10.

DeepSpeed

O Azure Machine Learning dá suporte ao DeepSpeed como um cidadão de primeira classe para executar trabalhos distribuídos com escalabilidade quase linear em termos de:

• Aumento no tamanho do modelo
• Aumento no número de GPUs

O DeepSpeed pode ser habilitado usando a distribuição do Pytorch ou a MPI para executar o treinamento distribuído. O Azure Machine Learning dá suporte ao inicializador do DeepSpeed para iniciar o treinamento distribuído e também ao ajuste automático para obter a configuração ds ideal.

Você pode usar um ambiente selecionado como um ambiente pronto para uso com as tecnologias de última geração mais recentes — incluindo o DeepSpeed, ORT, MSSCCL e Pytorch — para seus trabalhos de treinamento do DeepSpeed.

Exemplo do DeepSpeed

• Para obter exemplos de treinamento e ajuste automático do DeepSpeed, consulte estas pastas.

TensorFlow

Se estiver usando o TensorFlow distribuído nativo no código do seu treinamento, como a API tf.distribute.Strategy do TensorFlow 2.x, você pode inicializar o trabalho distribuído por meio do Azure Machine Learning usando os parâmetros de distribution do objeto TensorFlowDistribution.

# create the command
job = command(
    code="./src",  # local path where the code is stored
    command="python main.py --epochs ${{inputs.epochs}} --model-dir ${{inputs.model_dir}}",
    inputs={"epochs": 1, "model_dir": "outputs/keras-model"},
    environment="AzureML-tensorflow-2.16-cuda12@latest",
    compute="cpu-cluster",
    instance_count=2,
    # distribution = {"type": "mpi", "process_count_per_instance": 1},
    # distribution={
    #     "type": "tensorflow",
    #     "parameter_server_count": 1,  # for legacy TensorFlow 1.x
    #     "worker_count": 2,
    #     "added_property": 7,
    # },
    # distribution = {
    #        "type": "pytorch",
    #        "process_count_per_instance": 4,
    #        "additional_prop": {"nested_prop": 3},
    #    },
    display_name="tensorflow-mnist-distributed-example"
    # experiment_name: tensorflow-mnist-distributed-example
    # description: Train a basic neural network with TensorFlow on the MNIST dataset, distributed via TensorFlow.
)

# can also set the distribution in a separate step and using the typed objects instead of a dict
job.distribution = TensorFlowDistribution(worker_count=2)

Se seu script de treinamento usar a estratégia de servidor de parâmetros para o treinamento distribuído, como no caso do TensorFlow 1.x herdado, você também vai precisar especificar o número de servidores de parâmetros a serem usados no trabalho, dentro do parâmetro distribution do command. No exemplo acima, por exemplo, "parameter_server_count" : 1 e "worker_count": 2.

TF_CONFIG

No TensorFlow, a variável de ambiente TF_CONFIG é necessária para treinamento em vários computadores. Nos trabalhos do TensorFlow, o Azure Machine Learning configura e define adequadamente a variável TF_CONFIG para cada trabalho antes de executar seu script de treinamento.

Se for preciso, você pode acessar o TF_CONFIG a partir de seu script de treinamento: os.environ['TF_CONFIG'].

Exemplo de TF_CONFIG definido em um nó de trabalho principal:

TF_CONFIG='{
    "cluster": {
        "worker": ["host0:2222", "host1:2222"]
    },
    "task": {"type": "worker", "index": 0},
    "environment": "cloud"
}'

Exemplo do TensorFlow

• Para que o notebook completo execute o exemplo do TensorFlow, confira azureml-examples: Treinar uma rede neural básica com a MPI distribuída no conjunto de dados do MNIST usando o Tensorflow com Horovod.

Acelerar o treinamento de GPU com o InfiniBand

Conforme aumenta o número de VMs que treinam um modelo, o tempo necessário para treinar esse modelo diminuirá. A redução no tempo, idealmente, deve ser linearmente proporcional ao número de VMs de treinamento. Por exemplo, se treinar um modelo em uma VM leva 100 segundos, o treinamento do mesmo modelo em duas VMs idealmente deve levar 50 segundos. Treinar o modelo em quatro VMs deve levar 25 segundos e assim por diante.

A InfiniBand pode ser um fator importante na obtenção desse dimensionamento linear. A InfiniBand permite comunicação de GPU para GPU de baixa latência entre os nós em um cluster. O InfiniBand requer um hardware especializado para operar. Algumas séries de VMs do Azure, especificamente as séries NC, ND e H, agora têm VMs compatíveis com RDMA e com suporte a SR-IOV e InfiniBand. Essas VMs se comunicam pela rede InfiniBand de baixa latência e alta largura de banda, que é muito mais eficaz do que a conectividade baseada em Ethernet. O SR-IOV para InfiniBand habilita o desempenho quase bare-metal para qualquer biblioteca de MPI (o MPI é usado por muitas ferramentas e estruturas de treinamento distribuídas, incluindo o software NCCL da NVIDIA.) Esses SKUs se destinam a satisfazer as necessidades de cargas de trabalho de aprendizado de máquina aceleradas por GPU e com uso intensivo de computação. Para obter mais informações, confira Acelerar o treinamento distribuído no Azure Machine Learning com o SR-IOV.

De modo geral, os SKUs de VM com um "r" no nome contêm o hardware de InfiniBand necessário, e aqueles sem um "r" de modo geral não. ("r" é uma referência ao RDMA, que significa acesso direto remoto à memória.) Por exemplo, o SKU de VM Standard_NC24rs_v3 é habilitado para InfiniBand, mas o SKU Standard_NC24s_v3 não. Descontando os recursos da InfiniBand, as especificações desses dois SKUs são praticamente as mesmas. Ambos têm 24 núcleos, 448 GB de RAM, 4 GPUs com o mesmo SKU etc. Saiba mais sobre as SKUs de máquina habilitadas para RDMA e InfiniBand.

Se você criar um cluster AmlCompute de um desses tamanhos compatíveis com RDMA e habilitados para InfiniBand, a imagem do sistema operacional virá com o driver Mellanox OFED necessário para habilitar a InfiniBand pré-instalada e pré-configurada.

Próximas etapas

• Implantar e pontuar um modelo de aprendizado de máquina usando um ponto de extremidade online
• Arquitetura de referência para treinamento de aprendizado profundo distribuído no Azure
• Solução de problemas de ambiente