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

APLICA-SE A: SDK 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 os recursos de treinamento distribuído nativo do PyTorch (torch.distributed).

Inicialização do grupo de processos

O backbone de qualquer treinamento distribuído é um grupo de processos que se conhecem e podem se comunicar entre si usando um back-end. Para o PyTorch, você cria o grupo de processos chamando torch.distributed.init_process_group em todos os processos distribuídos para formar coletivamente um grupo de processos.

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

Os back-ends de comunicação mais comuns são mpi, nccle gloo. Para treinamento baseado em GPU, use nccl para obter o melhor desempenho.

O init_method parâmetro especifica como cada processo descobre os outros processos e como eles inicializam e verificam o grupo de processos usando o back-end de comunicação. Por padrão, se você não especificar init_method, o PyTorch usará o método de inicialização da variável de ambiente (env://). Use init_method em seu 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ó precisam ser executadas uma vez por nó, geralmente em 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 geralmente deve ser igual a # of GPUs per node. Se você não especificar process_count_per_instance, o Azure Machine Learning iniciará um processo por nó por padrão.

O Azure Machine Learning define as variáveis de ambiente MASTER_ADDR, MASTER_PORT, WORLD_SIZE e NODE_RANK em cada nó. Ele define as variáveis de ambiente em nível de processo de RANK e 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.8-cuda12.6@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_NC4AS_T4_V3", instance_count=2
)  # Serverless compute resources

Exemplo de PyTorch

• Para o notebook completo para executar o exemplo de PyTorch, consulte 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

Você pode habilitar o DeepSpeed usando a distribuição PyTorch ou o 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 coletado para um ambiente pronto para uso com as mais recentes tecnologias de ponta, incluindo DeepSpeed, ORT, MSSCCL e PyTorch para seus trabalhos de treinamento DeepSpeed.

Exemplo do DeepSpeed

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

TensorFlow

Se você usar TensorFlow nativo distribuído no código de treinamento, como a tf.distribute.Strategy API do TensorFlow 2.x, poderá iniciar o trabalho distribuído por meio do Azure Machine Learning usando parâmetros distribution ou o 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 anterior, você especifica "parameter_server_count" : 1 e "worker_count": 2.

TF_CONFIG

No TensorFlow, você precisa da variável de TF_CONFIG ambiente para treinar em vários computadores. Para trabalhos do TensorFlow, o Azure Machine Learning define a TF_CONFIG variável corretamente 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 tensorFlow, consulte azureml-examples: Treinar uma rede neural básica com MPI distribuída no conjunto de dados MNIST usando TensorFlow com Horovod.

Acelerar o treinamento de GPU com o InfiniBand

À medida que você aumenta o número de VMs treinando um modelo, o tempo necessário para treinar esse modelo deve diminuir. A diminuição no tempo deve ser linearmente proporcional ao número de máquinas virtuais de treinamento. Por exemplo, se o treinamento de um modelo em uma VM levar 100 segundos, o treinamento do mesmo modelo em duas VMs deverá levar 50 segundos. Treinar o modelo em quatro VMs deve levar 25 segundos e assim por diante.

O InfiniBand pode ajudá-lo a obter esse 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 permite um desempenho próximo ao bare-metal para qualquer biblioteca MPI (MPI é usado por muitas ferramentas e estruturas de treinamento distribuído, incluindo o software NCCL da NVIDIA). Essas SKUs destinam-se a atender às necessidades de cargas de trabalho de machine learning intensivas em computação e aceleradas por GPU. Para obter mais informações, confira Acelerar o treinamento distribuído no Azure Machine Learning com o SR-IOV.

Normalmente, SKUs de VM com um "r" em seu nome contêm o hardware InfiniBand necessário, e aqueles sem um "r" normalmente não contêm. ("r" é uma referência ao RDMA, que significa acesso remoto direto à memória.) Por exemplo, o SKU Standard_NC24rs_v3 da VM está habilitado para InfiniBand, mas o SKU Standard_NC24s_v3 não está. 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 da mesma SKU e assim por diante. 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 machine learning 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