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 CPUnccl
: 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 como0
. Em cada worker, ele é definido com um valor diferente de1
paraWORLD_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
- Para saber mais sobre o treinamento com o PyTorch no AI Platform Training, siga o tutorial Primeiros passos com o PyTorch.
- Para saber mais sobre o treinamento distribuído do PyTorch em geral, leia o Guia da documentação do PyTorch para treinamento distribuído.