Como configurar o treinamento distribuído do PyTorch

Neste documento, você verá como criar um job de treinamento distribuído do PyTorch. Ao criar um job de treinamento distribuído, o AI Platform Training executa o código em um cluster de instâncias de máquina virtual (VM, na sigla em inglês), também conhecido como nós, com variáveis de ambiente compatíveis com o treinamento distribuído do PyTorch. Isso ajuda o job de treinamento a escalonar para lidar com uma grande quantidade de dados.

Neste guia, presume-se que você esteja usando contêineres PyTorch predefinidos para treinamento, conforme descrito em Primeiros passos com o PyTorch. A adaptação do código PyTorch para treinamento distribuído requer alterações mínimas.

Como especificar a estrutura do cluster de treinamento

Para o treinamento distribuído do PyTorch, configure o job para usar um nó de worker mestre e um ou mais nós de worker. Esses papéis têm os seguintes comportamentos:

  • Worker mestre: a VM com classificação 0. Esse nó configura conexões entre os nós no cluster.
  • Worker: os nós restantes no cluster. Cada nó faz uma parte do treinamento, conforme especificado pelo código do aplicativo de treinamento.

Para saber como especificar o nó de worker mestre e os nós de worker do seu cluster de treinamento, leia Como especificar tipos de máquinas ou níveis de escalonamento.

Como especificar imagens de contêiner

Ao criar um job de treinamento, especifique a imagem de um contêiner do Docker para o worker mestre usar no campo trainingInput.masterConfig.imageUri e especifique a imagem de um contêiner do Docker para cada worker usar no campo trainingInput.workerConfig.imageUri. Veja a lista de contêineres PyTorch pré-criados.

Se você usar o comando gcloud ai-platform jobs submit training para criar o job de treinamento, especifique esses campos com as sinalizações --master-image-uri e --worker-image-uri.

No entanto, se você não especificar o campo trainingInput.workerConfig.imageUri, o valor padrão será o valor de trainingInput.masterConfig.imageUri. Geralmente, faz sentido usar o mesmo contêiner PyTorch pré-criado em cada nó.

Como atualizar o código de treinamento

No aplicativo de treinamento, adicione o seguinte código para inicializar o cluster de treinamento:

import torch

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

Substitua BACKEND por um dos back-ends de treinamento distribuídos compatíveis descritos na seção a seguir. O argumento de palavra-chave init_method='env://' diz ao PyTorch para usar variáveis de ambiente para inicializar a comunicação no cluster. Saiba mais na seção Variáveis de ambiente deste guia.

Além disso, atualize seu código de treinamento para usar a classe torch.nn.parallel.DistributedDataParallel. Por exemplo, se você tiver criado um módulo PyTorch chamado model no código, adicione a linha a seguir:

model = torch.nn.parallel.DistributedDataParallel(model)

Para saber mais sobre como configurar o treinamento distribuído, leia o guia da documentação do PyTorch para treinamento distribuído.

Back-ends de treinamento distribuído

O AI Platform Training é compatível com os seguintes back-ends para treinamento distribuído do PyTorch:

  • gloo: recomendado para jobs de treinamento de CPU
  • nccl: recomendado para jobs de treinamento da GPU

Leia sobre as diferenças entre back-ends.

Variáveis de ambiente

Quando você cria um job de treinamento distribuído do PyTorch, o AI Platform Training define as seguintes variáveis de ambiente em cada nó:

  • WORLD_SIZE: o número total de nós no cluster. Essa variável tem o mesmo valor em todos os nós.
  • RANK: um identificador exclusivo para cada nó. No worker mestre, isso é definido como 0. Em cada worker, ele é definido com um valor diferente de 1 para WORLD_SIZE - 1.
  • MASTER_ADDR: o nome do host do nó de worker mestre. Essa variável tem o mesmo valor em todos os nós.
  • MASTER_PORT: a porta em que o nó de worker mestre se comunica. Essa variável tem o mesmo valor em todos os nós.

O PyTorch usa essas variáveis de ambiente para inicializar o cluster.

A seguir