Treinamento distribuído

Nesta página, você aprenderá como executar jobs de treinamento distribuído na Vertex AI.

Exigências de código

usar um framework de ML compatível com o treinamento distribuído; No código de treinamento, use as variáveis de ambiente CLUSTER_SPEC ou TF_CONFIG para fazer referência a partes específicas do cluster de treinamento.

Estrutura do cluster de treinamento

Ao executar um job de treinamento distribuído com a Vertex AI, você especifica várias máquinas (nós) em um cluster de treinamento. Esse serviço aloca os recursos dos tipos de máquina que você especifica. O job em execução em um determinado nó é denominado réplica. Um grupo de réplicas com a mesma configuração é chamado de pool de workers.

Cada réplica no cluster de treinamento recebe um único papel ou tarefa no treinamento distribuído. Exemplo:

• Réplica principal: exatamente uma réplica é designada como réplica principal. Essa tarefa gerencia as outras e relata o status do job como um todo.

• Workers: uma ou mais réplicas podem ser designadas como workers. Elas fazem o trabalho conforme atribuído na configuração do job.

• Servidores de parâmetros: se compatível com seu framework de ML, uma ou mais réplicas podem ser designadas como servidores de parâmetros. Elas armazenam parâmetros do modelo e coordenam o estado do modelo compartilhado entre os workers.

• Avaliador(es): se compatível com seu framework de ML, uma ou mais réplicas poderão ser designadas como avaliadores. É possível usá-las para avaliar o modelo. Se você estiver usando o TensorFlow, observe que ele geralmente espera que você não use mais de um avaliador.

Configurar um job de treinamento distribuído

É possível configurar qualquer job de treinamento personalizado como um job de treinamento distribuído definindo vários pools de workers. Também é possível executar o treinamento distribuído em um pipeline de treinamento ou em um job de ajuste de hiperparâmetro.

Para configurar um job de treinamento distribuído, defina sua lista de pools de workers (workerPoolSpecs[]), criando um WorkerPoolSpec para cada tipo de tarefa:

Posição em workerPoolSpecs[]	Tarefa realizada no cluster
Primeira (workerPoolSpecs[0])	Principal, chefe, programador ou "mestre"
Segunda (workerPoolSpecs[1])	Secundário, réplicas, workers
Terceira (workerPoolSpecs[2])	Servidores de parâmetros, Reduction Server
Quarta (workerPoolSpecs[3])	Avaliadores

Especifique uma réplica principal, que coordena o trabalho feito por todas as outras réplicas. Use a primeira especificação do pool de workers apenas para sua réplica principal e defina a replicaCount como 1:

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, primary replica, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {},
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

Especificar outros pools de workers

Dependendo do framework de ML, é possível especificar mais pools de workers para outros fins. Por exemplo, se você estiver usando o TensorFlow, poderá especificar pools de workers para configurar réplicas de workers, réplicas de servidores de parâmetros e réplicas de avaliadores.

A ordem dos pools de workers especificados na lista workerPoolSpecs[] determina o tipo de pool de workers. Defina valores vazios para pools de workers que você não quer usar para que seja possível ignorá-los na lista workerPoolSpecs[] a fim de especificar os pools de workers que você quer usar. Exemplo:

Se quiser especificar um job que tenha apenas uma réplica principal e um pool de workers de servidor de parâmetros, defina um valor vazio para esse pool:

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

Reduza o tempo de treinamento com o servidor de redução

Ao treinar um modelo grande de ML usando vários nós, a comunicação de gradientes entre nós pode contribuir com uma latência significativa. O Reduction Server é um algoritmo totalmente reduzido que pode aumentar a capacidade e reduzir a latência do treinamento distribuído. O Vertex AI disponibiliza o Reduction Server em uma imagem de contêiner do Docker que pode ser usada para um dos pools de workers durante o treinamento distribuído.

Para saber como o Reduction Server funciona, consulte Treinamento de GPU distribuído mais rápido com o Reduction Server na Vertex AI.

Pré-requisitos

Use o Reduction Server se atender aos requisitos abaixo:

• Você está realizando treinamento distribuído com workers da GPU.

• O código de treinamento usa o TensorFlow ou PyTorch e está configurado para treinamento paralelo de dados de vários hosts com GPUs usando o NCCL totalmente reduzido. Também é possível usar outros frameworks de ML que usam o NCCL.

• Os contêineres em execução no nó principal (workerPoolSpecs[0]) e nos workers (workerPoolSpecs[1]) são compatíveis com o Reduction Server. Especificamente, cada contêiner é um dos seguintes:

  • Um contêiner de treinamento predefinido do TensorFlow, versão 2.3 ou posterior.

  • Um container de treinamento predefinido do Pytorch, versão 1.4 ou posterior.

  • Um container personalizado com o NCCL 2.7 ou posterior e o pacote google-reduction-server instalado. É possível instalar este pacote em uma imagem de contêiner personalizada adicionando a seguinte linha ao Dockerfile:

    RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list && \
        curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
        apt update && apt install -y google-reduction-server
    

Treinar usando o Reduction Server

Para usar o Reduction Server, faça o seguinte ao criar um recurso de treinamento personalizado:

1. Especifique um dos seguintes URIs no campo containerSpec.imageUri do terceiro pool de workers (workerPoolSpecs[2]):

   • us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
   • europe-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
   • asia-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

   Escolher a multirregião mais próxima de onde você está realizando treinamento personalizado pode reduzir a latência.

2. Ao selecionar o tipo de máquina e o número de nós do terceiro pool de workers, verifique se a largura de banda total da rede do terceiro pool de workers corresponde ou excede a largura de banda total da rede do primeiro e do segundo pools de workers.

   Para saber mais sobre a largura de banda máxima disponível de cada nó no segundo pool de workers, consulte Largura de banda de rede e GPUs.

   Não são usados GPUs para os nós do Reduction Server. Para saber mais sobre a largura de banda máxima disponível de cada nó no terceiro pool de workers, consulte as colunas "Largura de banda máxima de saída (Gbps)" em Família de máquinas de uso geral.

   Por exemplo, se você configurar o primeiro e o segundo pool de workers para usar cinco nós n1-highmem-96, cada um com oito GPUs NVIDIA_TESLA_V100, cada nó terá uma largura de banda máxima disponível de 100 Gbps para uma largura de banda total de 500 Gbps. Para corresponder a essa largura de banda no terceiro pool de workers, use 16 nós n1-highcpu-16, cada um com uma largura de banda máxima de 32 Gbps para uma largura de banda total de 512 Gbps.

   Recomendamos que você use o tipo de máquina n1-highcpu-16 para nós do Reduction Server, porque esse tipo de máquina oferece uma largura de banda relativamente alta para os respectivos recursos.

O comando a seguir fornece um exemplo de como criar um recurso CustomJob que usa o Reduction Server:

gcloud ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=4,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highcpu-16,replica-count=16,container-image-uri=us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

Para mais contexto, consulte o guia de criação de um CustomJob.

Para acessar um exemplo de como realizar treinamento distribuído no PyTorch usando o Reduction Server, execute o notebook do Jupyter "Treinamento distribuído do PyTorch com o servidor de redução da Vertex AI" em um dos seguintes ambientes:

Abrir no Colab | Abrir no Colab Enterprise | Abrir nos blocos gerenciados pelo usuário do Vertex AI Workbench | Ver no GitHub

Para ver um exemplo de como realizar o treinamento paralelo no XGBoost usando o Dask, execute o notebook Jupyter "Distributed XGBoost training with Dask" em um dos seguintes ambientes:

Abrir no Colab | Abrir no Colab Enterprise | Abrir nos blocos gerenciados pelo usuário do Vertex AI Workbench | Ver no GitHub

Práticas recomendadas para treinamento usando o Reduction Server

Tipo e contagem de máquinas

No treinamento do Reduction Server, cada worker precisa se conectar a todos os hosts redutores. Para minimizar o número de conexões no host do worker, use um tipo de máquina com a maior largura de banda de rede para o host redutor.

Uma boa opção para hosts redutores é uma VM de uso geral N1/N2 com pelo menos 16 vCPU que fornece uma largura de banda de saída de 32 Gbps, como n1-highcpu-16 e n2-highcpu-16. A largura de banda da VM de nível 1 para VMs N1/N2 aumenta a largura de banda de saída máxima entre 50 Gbps e 100 Gbps, o que torna essa uma boa opção para nós de VM redutores.

A largura de banda de saída total de workers e redutores precisa ser igual. Por exemplo, se você usar oito VMs a2-megagpu-16g como workers, use pelo menos 25 VMs n1-highcpu-16 como redutores.

`(8 worker VMs * 100 Gbps) / 32 Gbps egress = 25 reducer VMs`.

Agrupar mensagens pequenas em lote

O Reduction Server vai funcionar melhor se as mensagens a serem agregadas forem grandes o suficiente. A maioria dos frameworks de ML já fornece técnicas com terminologia diferente para agrupar tensores de gradiente pequeno antes de executar o all-reduce.

Horovod

O Horovod é compatível com o Tensor Fusion para agrupar pequenos tensores para o all-reduce. Os tensores são preenchidos em um buffer de fusão até que o buffer esteja totalmente preenchido e a operação de all-reduce no buffer seja executada. É possível ajustar o tamanho do buffer de fusão definindo a variável de ambiente HOROVOD_FUSION_THRESHOLD.

O valor recomendado para a variável de ambiente HOROVOD_FUSION_THRESHOLD é de pelo menos 128 MB. Nesse caso, defina a variável de ambiente HOROVOD_FUSION_THRESHOLD como 134217728 (128 * 1024 * 1024).

PyTorch

O DistributedDataParallel do PyTorch oferece suporte a mensagens em lote como "agrupamento por gradiente". Defina o parâmetro bucket_cap_mb no construtor DistributedDataParallel para controlar o tamanho dos buckets em lote. O tamanho padrão é 25 MB.

PRÁTICA RECOMENDADA: o valor recomendado de bucket_cap_mb é 64 (64 MB).

Variáveis de ambiente para o cluster

A Vertex AI preenche uma variável de ambiente CLUSTER_SPEC em cada réplica para descrever como o cluster geral está configurado. Como o TF_CONFIG do TensorFlow, o CLUSTER_SPEC descreve todas as réplicas no cluster, incluindo o índice e o papel (worker mestre, worker, servidor de parâmetros ou avaliador).

Ao executar o treinamento distribuído com o TensorFlow, o TF_CONFIG é analisado para criar tf.train.ClusterSpec. Da mesma maneira, ao executar o treinamento distribuído com outros frameworks de ML, você precisa analisar CLUSTER_SPEC para preencher as variáveis de ambiente ou as configurações exigidas pelo framework.

Formato de CLUSTER_SPEC

A variável de ambiente CLUSTER_SPEC é uma string JSON com o seguinte formato:

Chave	Descrição
"cluster"	A descrição do cluster referente ao contêiner personalizado. Assim como acontece com o TF_CONFIG, esse objeto é formatado como uma especificação de cluster do TensorFlow e é possível transmiti-lo para o construtor do tf.train.ClusterSpec. A descrição do cluster contém uma lista de nomes de réplica para cada pool de worker que você especificar.
	"workerpool0"	Todos os jobs de treinamento distribuídos têm uma réplica principal no primeiro pool de workers.
	"workerpool1"	Este pool de workers contém réplicas de workers, se você as tiver especificado ao criar o job.
	"workerpool2"	Este pool de workers contém servidores de parâmetros, se você os tiver especificado ao criar o job.
	"workerpool3"	Esse pool de workers contém avaliadores, se você os especificou ao criar o job.
"environment"	A string cloud.
"task"	Descreve a tarefa do nó específico em que o código está sendo executado. Você pode usar essas informações para escrever código para trabalhadores específicos de um job distribuído. Esta entrada é um dicionário com as seguintes chaves:
	"type"	O tipo de pool de worker em que esta tarefa está sendo executada. Por exemplo, "workerpool0" refere-se à réplica primária.
	"index"	O índice de tarefa com base em zero. Por exemplo, se seu job de treinamento incluir dois workers, este valor será definido como 0 em um deles e 1 no outro.
	"trial"	O identificador do teste de ajuste de hiperparâmetros em execução no momento. Ao configurar o ajuste de hiperparâmetros para o job, você define uma série de testes para treinamento. Com esse valor, você diferencia os testes que estão sendo executados no código. O identificador é um valor de string contendo o número do teste, começando com 1.
job	O CustomJobSpec que você forneceu para criar o job de treinamento atual, representado como um dicionário.

CLUSTER_SPEC exemplo

Veja um exemplo de valor:

{
   "cluster":{
      "workerpool0":[
         "cmle-training-workerpool0-ab-0:2222"
      ],
      "workerpool1":[
         "cmle-training-workerpool1-ab-0:2222",
         "cmle-training-workerpool1-ab-1:2222"
      ],
      "workerpool2":[
         "cmle-training-workerpool2-ab-0:2222",
         "cmle-training-workerpool2-ab-1:2222"
      ],
      "workerpool3":[
         "cmle-training-workerpool3-ab-0:2222",
         "cmle-training-workerpool3-ab-1:2222",
         "cmle-training-workerpool3-ab-2:2222"
      ]
   },
   "environment":"cloud",
   "task":{
      "type":"workerpool0",
      "index":0,
      "trial":"TRIAL_ID"
   },
   "job": {
      ...
   }
}

Formato de TF_CONFIG

Além de CLUSTER_SPEC, a Vertex AI define a variável de ambiente TF_CONFIG em cada réplica de todos os jobs de treinamento distribuídos. A Vertex AI não define TF_CONFIG para jobs de treinamento de réplica única.

CLUSTER_SPEC e TF_CONFIG compartilham alguns valores, mas têm formatos diferentes. As duas variáveis de ambiente incluem campos adicionais além do que o TensorFlow requer.

O treinamento distribuído com o TensorFlow funciona da mesma maneira quando você usa contêineres personalizados, como ao usar um contêiner pré-criado.

A variável de ambiente TF_CONFIG é uma string JSON com o seguinte formato:

Campos de TF_CONFIG
cluster	Descrição do cluster do TensorFlow. Um dicionário que correlaciona um ou mais nomes de tarefa (chief, worker, ps ou master) com listas de endereços de rede em que essas tarefas estão em execução. Para um determinado job de treinamento, esse dicionário é o mesmo em todas as VMs. Este é um primeiro argumento válido para o construtor tf.train.ClusterSpec. Observe que este dicionário nunca contém evaluator como chave, já que os avaliadores não são considerados parte do cluster de treinamento, mesmo se você usá-los para o job.
task	A descrição da tarefa da VM em que esta variável de ambiente está definida. Para um determinado job de treinamento, esse dicionário é diferente em cada VM. É possível usar essas informações para personalizar o código executado em cada VM em um job de treinamento distribuído. Também é possível usá-las para alterar o comportamento de seu código de treinamento para diferentes testes de um job de ajuste de hiperparâmetros. Esse dicionário inclui os seguintes pares de chave-valor: Campos task type O tipo de tarefa que a VM está executando. Esse valor é definido como worker nos workers, ps nos servidores de parâmetros e evaluator nos avaliadores. No worker mestre do job, o valor é definido como chief ou master. Saiba mais sobre a diferença entre os dois na seção chief x master deste documento. index O índice de tarefa com base em zero. Por exemplo, se seu job de treinamento incluir dois workers, este valor será definido como 0 em um deles e 1 no outro. trial O ID do teste de ajuste de hiperparâmetros atualmente em execução nesta VM. Este campo só será definido se o job de treinamento atual for um job de ajuste de hiperparâmetros. Para jobs de ajuste de hiperparâmetros, a Vertex AI executa o código de treinamento repetidamente em vários testes com hiperparâmetros diferentes a cada vez. Este campo contém o número do teste atual, começando com 1 para o primeiro teste. cloud Um código usado internamente pela Vertex AI. É possível ignorar este campo.
job	O CustomJobSpec que você forneceu para criar o job de treinamento atual, representado como um dicionário.
environment	A string cloud.

TF_CONFIG exemplo

No código de exemplo a seguir, a variável de ambiente TF_CONFIG é impressa nos seus registros de treinamento:

import json
import os

tf_config_str = os.environ.get('TF_CONFIG')
tf_config_dict  = json.loads(tf_config_str)

# Convert back to string just for pretty printing
print(json.dumps(tf_config_dict, indent=2))

Em um job de ajuste de hiperparâmetros executado no ambiente de execução versão 2.1 ou posterior e que usa um worker mestre, dois workers e um servidor de parâmetros, este código produz o seguinte registro para um dos workers durante o primeiro teste de ajuste de hiperparâmetros. O exemplo de saída oculta o campo job para concisão e substitui alguns IDs por valores genéricos.

{
  "cluster": {
    "chief": [
      "training-workerpool0-[ID_STRING_1]-0:2222"
    ],
    "ps": [
      "training-workerpool2-[ID_STRING_1]-0:2222"
    ],
    "worker": [
      "training-workerpool1-[ID_STRING_1]-0:2222",
      "training-workerpool1-[ID_STRING_1]-1:2222"
    ]
  },
  "environment": "cloud",
  "job": {
    ...
  },
  "task": {
    "cloud": "[ID_STRING_2]",
    "index": 0,
    "trial": "1",
    "type": "worker"
  }
}

Quando usar TF_CONFIG

TF_CONFIG é definido apenas para jobs de treinamento distribuídos.

Você provavelmente não precisa interagir com a variável de ambiente TF_CONFIG diretamente em seu código de treinamento. Acesse a variável de ambiente TF_CONFIG apenas se as estratégias de distribuição do TensorFlow e o fluxo de trabalho de ajuste de hiperparâmetros padrão da Vertex AI, ambos descritos nas próximas seções, não funcionarem para seu job.

Treinamento distribuído

A Vertex AI define a variável de ambiente TF_CONFIG para ampliar as especificações necessárias para o treinamento distribuído pelo TensorFlow.

Para realizar treinamento distribuído com o TensorFlow, use a API tf.distribute.Strategy (em inglês). Em particular, recomendamos que você use a API Keras junto com o MultiWorkerMirroredStrategy ou, se você especificar servidores de parâmetros para seu job, o ParameterServerStrategy. No entanto, o TensorFlow atualmente só fornece compatibilidade experimental para essas estratégias.

Essas estratégias de distribuição usam a variável de ambiente TF_CONFIG para atribuir papéis a cada VM em seu job de treinamento e para facilitar a comunicação entre as VMs. Você não precisa acessar a variável de ambiente TF_CONFIG diretamente no seu código de treinamento, porque o TensorFlow cuida disso para você.

Somente analise a variável de ambiente TF_CONFIG diretamente se quiser personalizar o comportamento das diferentes VMs que executam seu job de treinamento.

Ajuste de hiperparâmetros

Quando você executar um job de ajuste de hiperparâmetros, a Vertex AI fornece argumentos diferentes para o código de treinamento de cada avaliação. Seu código de treinamento não precisa necessariamente estar ciente de qual teste está sendo realizado. Além disso, é possível monitorar o progresso dos jobs de ajuste de hiperparâmetros no console do Google Cloud.

Se necessário, seu código poderá ler o número do teste atual no campo trial do campo task da variável de ambiente TF_CONFIG.

A seguir

• Crie um pipeline de treinamento.