Treinamento distribuído com contêineres

Nesta página, você verá como os contêineres personalizados são compatíveis com a estrutura do treinamento distribuído no AI Platform Training.

Com os contêineres personalizados, você faz o treinamento distribuído com qualquer framework de ML que aceite a distribuição. A terminologia usada aqui é baseada no modelo distribuído do TensorFlow. No entanto, é possível usar qualquer outro framework de ML que tenha uma estrutura de distribuição semelhante. Por exemplo, o treinamento distribuído no MXNet (em inglês) usa um programador, workers e servidores. Isso corresponde à estrutura de treinamento distribuída de contêineres personalizados do AI Platform Training, que usa um mestre, workers e servidores de parâmetro.

Estrutura do cluster de treinamento

Se você executar um job de treinamento distribuído com o AI Platform Training, especifique várias máquinas (nós) em um cluster de treinamento. Esse serviço alocará os recursos para os tipos de máquina que você especificar. O job em execução em um determinado nó é denominado réplica. De acordo com o modelo distribuído do TensorFlow, cada réplica no cluster de treinamento recebe um único papel ou tarefa no treinamento distribuído:

  • Worker mestre: exatamente uma réplica é designada como worker mestre (também conhecida como worker chief). Essa tarefa gerencia as outras e relata o status do job como um todo.

  • Worker(s): 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: 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): uma ou mais réplicas podem 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.

Mapeamento da API

É possível atribuir quatro papéis diferentes a máquinas no cluster de treinamento. Eles correspondem a quatro campos que você especifica em TrainingInput, que representa os parâmetros de entrada para um job de treinamento:

  • masterConfig.imageUri representa o URI da imagem do contêiner a ser executado no worker mestre.
  • workerConfig.imageUri, parameterServerConfig.imageUri e evaluatorConfig.imageUri representam os URIs da imagem do contêiner que serão executados nos workers, servidores de parâmetros e avaliadores, respectivamente. Se nenhum valor for definido para esses campos, o AI Platform Training usará masterConfig.imageUri.

É possível também definir os valores para cada um desses campos com suas sinalizações correspondentes em gcloud ai-platform jobs submit training:

  • Para a configuração do worker mestre, use --master-image-uri.
  • Para a configuração do worker, use --worker-image-uri.
  • Para a configuração do servidor de parâmetros, use --parameter-server-image-uri.
  • No momento, não há uma sinalização para especificar o URI da imagem do contêiner para os avaliadores. É possível especificar evaluatorConfig.imageUri em um arquivo de configuração config.yaml.

Veja um exemplo de como enviar um job de treinamento distribuído com contêineres personalizados.

Noções básicas sobre CLUSTER_SPEC

O AI Platform Training preenche a 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 (em inglês).
"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 tarefa desempenhada pelo nó. Os valores possíveis são master, worker, ps e evaluator.
"index" O índice de tarefa com base em zero. A maioria dos jobs de treinamento distribuído tem uma tarefa mestre única, um ou mais servidores de parâmetros e um ou mais trabalhadores.
"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" Os parâmetros do job usado quando ele foi iniciado. Na maioria dos casos, é possível ignorar essa entrada, porque ela replica os dados transmitidos para o aplicativo por meio dos argumentos de linha de comando dele.

Comparação com o TF_CONFIG

Observe que o AI Platform Training também define a variável de ambiente TF_CONFIG em cada réplica de todos os jobs de treinamento. O AI Platform Training define apenas a variável de ambiente CLUSTER_SPEC em réplicas de jobs de treinamento de contêineres personalizados. As duas variáveis de ambiente compartilham alguns valores, mas têm formatos diferentes.

Ao treinar com contêineres personalizados, a réplica mestre será rotulada na variável de ambiente TF_CONFIG com o nome de tarefa master por padrão. É possível configurá-la para ser rotulada com o nome de tarefa chief configurando o campo trainingInput.useChiefInTfConfig como true ao criar o job de treinamento ou usando uma ou mais réplicas do avaliador no job. Isso é especialmente útil se o contêiner personalizado usar o TensorFlow 2.

Além dessa opção de configuração, o treinamento distribuído com o TensorFlow funciona da mesma maneira usando contêineres personalizados ou o ambiente de execução do AI Platform Training. Veja mais detalhes e exemplos de como usar o TF_CONFIG em treinamentos distribuídos no AI Platform Training.

A seguir