Como usar GPUs para treinar modelos na nuvem

As unidades de processamento de gráfico (GPUs, na sigla em inglês) aceleram significativamente o processo de treinamento para muitos modelos de aprendizado profundo. Modelos de treinamento para tarefas como classificação de imagem, análise de vídeo e processamento de linguagem natural envolvem multiplicação de matriz com uso intensivo de computação e outras operações que podem aproveitar a arquitetura massivamente paralela de uma GPU.

Leva dias para executar um treinamento de modelo de aprendizado profundo que envolve tarefas de computação intensivas em conjuntos de dados extremamente grandes em um único processador. No entanto, projetar o programa para descarregar essas tarefas em uma ou mais GPUs reduz o tempo de treinamento para horas, em vez de dias.

Antes de começar

Com o AI Platform Training, é possível executar seu aplicativo de treinamento do TensorFlow em uma máquina ativada para GPU. Leia o guia do TensorFlow para usar GPUs (em inglês) e a seção deste documento sobre como ajustar o código de treinamento para utilizar GPUs para garantir que seu aplicativo use as GPUs disponíveis.

Também é possível usar GPUs com frameworks de machine learning que não são do TensorFlow, basta usar um contêiner personalizado para treinamento.

Alguns modelos não são beneficiados com a execução em GPUs. Recomendamos GPUs para modelos grandes e complexos com muitas operações matemáticas. Mesmo assim, é aconselhável testar o benefício da compatibilidade de GPU executando uma pequena amostra dos dados por meio de treinamento.

Como solicitar máquinas ativadas para GPU

Para usar GPUs na nuvem, configure o job de treinamento para acessar máquinas ativadas para GPU de uma das seguintes maneiras:

  • usando o nível de escalonamento BASIC_GPU;
  • usando os tipos de máquina do Compute Engine e anexando as GPUs;
  • usando os tipos de máquina legados ativados para GPU.

Máquina básica ativada para GPU

Se você estiver aprendendo a usar o AI Platform Training ou testando máquinas ativadas para GPU, defina o nível de escalonamento como BASIC_GPU para acessar uma única instância de worker com uma GPU.

Tipos de máquina do Compute Engine com anexos de GPU

Se você configurar o job de treinamento com tipos de máquina do Compute Engine, será possível anexar um número personalizado de GPUs para acelerar o job:

  • Defina o nível de escalonamento como CUSTOM.
  • Configure o worker mestre e qualquer outro tipo de tarefa (worker, servidor de parâmetros ou avaliador) que faça parte do job para usar tipos de máquina válidos do Compute Engine.
  • Adicione um campo acceleratorConfig com o tipo e o número de GPUs que quiser a masterConfig, workerConfig, parameterServerConfig ou evaluatorConfig, dependendo das instâncias de máquina virtual (VM) que você quer acelerar. Também é possível usar os seguintes tipos de GPU:
    • NVIDIA_TESLA_A100
    • NVIDIA_TESLA_P4
    • NVIDIA_TESLA_P100
    • NVIDIA_TESLA_T4
    • NVIDIA_TESLA_V100

Para criar um acceleratorConfig válido, considere várias restrições:

  1. Só é possível usar determinados números de GPUs na configuração. Por exemplo, você pode anexar duas ou quatro NVIDIA Tesla T4s, mas não três. Para ver quais contagens são válidas para cada tipo de GPU, consulte a tabela de compatibilidade abaixo.

  2. É preciso garantir que cada uma das configurações de GPU forneça CPUs virtuais e memória suficientes para o tipo de máquina a que são anexadas. Por exemplo, se você usar n1-standard-32 para seus workers, cada worker terá 32 CPUs virtuais e 120 GB de memória. Como cada NVIDIA Tesla V100 pode fornecer até 12 CPUs virtuais e 76 GB de memória, é preciso anexar pelo menos quatro a cada worker n1-standard-32 para atender aos requisitos. Duas GPUs fornecem recursos insuficientes, e não é possível especificar três GPUs.

    Analisar a lista de tipos de máquina para o AI Platform Training e a comparação entre GPUs para cargas de trabalho de computação para determinar essas compatibilidades; ou consulte a tabela de compatibilidade abaixo.

    Observe a seguinte limitação adicional nos recursos de GPU para o AI Platform Training em casos específicos:

    • Uma configuração com quatro GPUs NVIDIA Tesla P100 permite até 64 CPUs virtuais e até 208 GB de memória em todas as regiões e zonas.
  3. É preciso enviar seu trabalho de treinamento para uma região compatível com sua configuração de GPU. Leia sobre a compatibilidade com regiões abaixo.

A tabela a seguir apresenta uma referência rápida de quantos tipos de acelerador é possível anexar a cada tipo de máquina do Compute Engine:

Números válidos de GPUs para cada tipo de máquina
Tipo de máquina NVIDIA A100 NVIDIA Tesla K80 NVIDIA Tesla P4 NVIDIA Tesla P100 NVIDIA Tesla T4 NVIDIA Tesla V100
n1-standard-4 1, 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 1, 2, 4, 8
n1-standard-8 1, 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 1, 2, 4, 8
n1-standard-16 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 2, 4, 8
n1-standard-32 4, 8 2, 4 2, 4 2, 4 4, 8
n1-standard-64 4 4 8
n1-standard-96 4 4 8
n1-highmem-2 1, 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 1, 2, 4, 8
n1-highmem-4 1, 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 1, 2, 4, 8
n1-highmem-8 1, 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 1, 2, 4, 8
n1-highmem-16 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 2, 4, 8
n1-highmem-32 4, 8 2, 4 2, 4 2, 4 4, 8
n1-highmem-64 4 4 8
n1-highmem-96 4 4 8
n1-highcpu-16 2, 4, 8 1, 2, 4 1, 2, 4 1, 2, 4 2, 4, 8
n1-highcpu-32 4, 8 2, 4 2, 4 2, 4 4, 8
n1-highcpu-64 8 4 4 4 8
n1-highcpu-96 4 4 8
a2-highgpu-1g 1
a2-highgpu-2g 2
a2-highgpu-4g 4
a2-highgpu-8g 8
a2-megagpu-16g 16

Veja abaixo um exemplo de envio de job usando os tipos de máquina do Compute Engine com as GPUs anexadas.

Tipos de máquina com GPUs incluídas

Como alternativa, em vez de usar um acceleratorConfig, selecione um tipo de máquina legado que tenha GPUs incluídas:

  • Defina o nível de escalonamento como CUSTOM.
  • Configure seu worker mestre e qualquer outro tipo de tarefa (worker, servidor de parâmetros ou avaliador) que você queira acelerar para usar um dos seguintes tipos de máquina ativados para GPU, com base no número de GPUs e o tipo de acelerador necessário para sua tarefa:
    • standard_gpu: uma única GPU
    • complex_model_m_gpu: quatro GPUs
    • complex_model_l_gpu: oito GPUs
    • standard_p100: uma única GPU NVIDIA Tesla P100
    • complex_model_m_p100: quatro GPUs NVIDIA Tesla P100
    • standard_v100: uma única GPU NVIDIA Tesla V100
    • large_model_v100: uma única GPU NVIDIA Tesla V100
    • complex_model_m_v100: quatro GPUs NVIDIA Tesla V100
    • complex_model_l_v100: oito GPUs NVIDIA Tesla V100

Veja abaixo um exemplo de envio de job com tipos de máquina ativados para GPU usando o comando gcloud.

Veja mais informações sobre tipos de máquinas para o AI Platform Training (em inglês).

Regiões compatíveis com GPUs

É necessário executar seu trabalho em uma região compatível com GPUs. No momento, as regiões que permitem acesso são:

  • us-west1
  • us-west2
  • us-central1
  • us-east1
  • us-east4
  • northamerica-northeast1
  • southamerica-east1
  • europe-west1
  • europe-west2
  • europe-west4
  • asia-south1
  • asia-southeast1
  • asia-east1
  • asia-northeast1
  • asia-northeast3
  • australia-southeast1

Além disso, algumas dessas regiões permitem acesso somente a determinados tipos de GPU. Para entender completamente as regiões disponíveis para os serviços do AI Platform Training, incluindo treinamento de modelo e previsão on-line/em lote, leia o guia de regiões.

Se o job de treinamento usa vários tipos de GPU, todos eles precisam estar disponíveis em uma única zona na região. Por exemplo, não é possível executar um job em us-central1 com um worker mestre usando GPUs NVIDIA Tesla T4, servidores de parâmetros usando GPUs NVIDIA Tesla K80 e workers usando GPUs NVIDIA Tesla P100. Todas essas GPUs estão disponíveis para jobs de treinamento em us-central1, mas nenhuma zona nessa região oferece os três tipos de GPU. Para saber mais sobre a disponibilidade de zonas das GPUs, confira a comparação entre GPUs para cargas de trabalho de computação.

Enviar o job de treinamento

É possível enviar o job de treinamento com o comando gcloud ai-platform jobs submit training.

  1. Defina um arquivo config.yaml que descreve as opções de GPU que você quer. A estrutura do arquivo YAML representa o recurso Job. Veja abaixo dois exemplos de arquivos config.yaml.

    No primeiro exemplo, mostramos um arquivo de configuração de um job de treinamento que usa os tipos de máquina do Compute Engine, alguns dos quais têm GPUs anexadas:

    trainingInput:
      scaleTier: CUSTOM
      # Configure a master worker with 4 T4 GPUs
      masterType: n1-highcpu-16
      masterConfig:
        acceleratorConfig:
          count: 4
          type: NVIDIA_TESLA_T4
      # Configure 9 workers, each with 4 T4 GPUs
      workerCount: 9
      workerType: n1-highcpu-16
      workerConfig:
        acceleratorConfig:
          count: 4
          type: NVIDIA_TESLA_T4
      # Configure 3 parameter servers with no GPUs
      parameterServerCount: 3
      parameterServerType: n1-highmem-8
    

    No próximo exemplo, mostramos um arquivo de configuração para um job com uma configuração semelhante à mostrada acima. No entanto, essa configuração usa tipos de máquinas legadas que incluem GPUs, em vez de anexar GPUs com acceleratorConfig:

    trainingInput:
      scaleTier: CUSTOM
      # Configure a master worker with 4 GPUs
      masterType: complex_model_m_gpu
      # Configure 9 workers, each with 4 GPUs
      workerCount: 9
      workerType: complex_model_m_gpu
      # Configure 3 parameter servers with no GPUs
      parameterServerCount: 3
      parameterServerType: large_model
    
  2. Use o comando gcloud para enviar o job, incluindo um argumento --config que aponta para o arquivo config.yaml. No exemplo a seguir, você precisa configurar as variáveis de ambiente para os valores de alguns argumentos. Elas são indicadas pelo sinal $ seguido de letras maiúsculas:

    gcloud ai-platform jobs submit training $JOB_NAME \
            --package-path $APP_PACKAGE_PATH \
            --module-name $MAIN_APP_MODULE \
            --job-dir $JOB_DIR \
            --region us-central1 \
            --config config.yaml \
            -- \
            --user_arg_1 value_1 \
             ...
            --user_arg_n value_n
    

Como alternativa, especifique detalhes de configuração de cluster com sinalizações de linha de comando, em vez de usar um arquivo de configuração. Saiba mais sobre como usar essas sinalizações.

Veja no exemplo a seguir como enviar um job com a mesma configuração do primeiro exemplo (usando os tipos de máquina do Compute Engine com GPUs anexadas), mas sem usar um arquivo config.yaml:

gcloud ai-platform jobs submit training $JOB_NAME \
        --package-path $APP_PACKAGE_PATH \
        --module-name $MAIN_APP_MODULE \
        --job-dir $JOB_DIR \
        --region us-central1 \
        --scale-tier custom \
        --master-machine-type n1-highcpu-16 \
        --master-accelerator count=4,type=nvidia-tesla-t4 \
        --worker-count 9 \
        --worker-machine-type n1-highcpu-16 \
        --worker-accelerator count=4,type=nvidia-tesla-t4 \
        --parameter-server-count 3 \
        --parameter-server-machine-type n1-highmem-8 \
        -- \
        --user_arg_1 value_1 \
         ...
        --user_arg_n value_n

Observações:

  • Caso especifique uma opção no arquivo de configuração (config.yaml) e como uma sinalização de linha de comando, o valor na linha de comando modificará o valor no arquivo de configuração.
  • A sinalização -- vazia marca o fim das sinalizações específicas do gcloud e o início do USER_ARGS que você quer passar para seu aplicativo.
  • Sinalizações específicas do AI Platform Training, como --module-name, --runtime-version e --job-dir, precisam ser inseridas antes da sinalização -- vazia. O serviço do AI Platform Training interpreta essas sinalizações.
  • É necessário que a sinalização --job-dir, se especificada, venha antes da sinalização vazia -- porque o AI Platform Training usa --job-dir para validar o caminho.
  • Seu aplicativo também precisa manipular a sinalização --job-dir, se especificada. Mesmo que a sinalização venha antes do -- vazio, --job-dir também será passado ao seu aplicativo como uma sinalização de linha de comando.
  • Defina quantos USER_ARGS forem necessários. O AI Platform Training passa --user_first_arg, --user_second_arg e assim por diante até seu aplicativo.

Para mais detalhes sobre as opções de envio de jobs, consulte o guia sobre como iniciar um job de treinamento.

Como ajustar o código de treinamento para usar GPUs

Se você usa Keras ou Estimators (links em inglês) para seu job de treinamento do TensorFlow e quer treinar usando uma única VM com uma GPU, não é necessário personalizar seu código para a GPU.

Se o cluster de treinamento tiver várias GPUs, use a API tf.distribute.Strategy no código de treinamento:

Para personalizar a forma como o TensorFlow atribui operações específicas às GPUs, leia o guia do TensorFlow para usar GPUs (em inglês). Neste caso, também pode ser útil saber como o AI Platform Training define a variável de ambiente TF_CONFIG em cada VM.

Strings do dispositivo da GPU

A única GPU da máquina standard_gpu é identificada como "/gpu:0". As máquinas com várias GPUs usam identificadores que começam com "/gpu:0", depois "/gpu:1" e assim por diante. Por exemplo, as máquinas complex_model_m_gpu contam com quatro GPUs identificadas como "/gpu:0" até "/gpu:3".

Pacotes Python em máquinas ativadas para GPU

As máquinas ativadas para GPU vêm com o tensorflow-gpu (em inglês) pré-instalado. Ele é o pacote do Python para TensorFlow compatível com GPUs. Consulte a Lista de versões do ambiente de execução para ver uma lista de todos os pacotes pré-instalados.

Eventos de manutenção

As VMs ativadas para GPU que executam jobs do AI Platform Training, às vezes, ficam sujeitas à manutenção do host do Compute Engine. As VMs são configuradas para serem reiniciadas automaticamente após esses eventos de manutenção, mas talvez seja necessário fazer algum trabalho extra para garantir que o job seja resiliente a essas paralisações. Configure seu aplicativo de treinamento para salvar regularmente checkpoints de modelo (geralmente no caminho do Cloud Storage especificado pelo argumento --job-dir para gcloud ai-platform jobs submit training) e restaurar o checkpoint mais recente, caso já exista um.

Os TensorFlow Estimators implementam essa funcionalidade para você, desde que você especifique um model_dir. Os Estimators salvam checkpoints regularmente para model_dir e tentam carregar do checkpoint mais recente. Portanto, não é preciso se preocupar com eventos de manutenção nos workers da GPU.

Se você estiver treinando com Keras, use o callback ModelCheckpoint (em inglês) para salvar regularmente o progresso do treinamento. Se você estiver usando tf.distribute.Strategy com Keras, suas VMs usarão checkpoints para se recuperar automaticamente das reinicializações (em inglês). Caso contrário, adicione lógica ao seu código de treinamento para verificar a existência de um checkpoint recente e restaurar do checkpoint, se ele existir (links em inglês).

Para casos mais avançados, leia o guia do TensorFlow para checkpoints (em inglês).

A seguir