Como treinar modelos PyTorch em pods do Cloud TPU

Neste tutorial, você verá como escalonar verticalmente o treinamento do modelo a partir de uma única Cloud TPU (v2-8 ou v3-8) para um pod da Cloud TPU. Os aceleradores das Cloud TPUs em um pod de TPU são conectados por uma interconexão de largura de banda muito alta. Isso os torna excelentes para escalonar verticalmente jobs de treinamento.

Para mais informações sobre as ofertas de pods da Cloud TPU, consulte a página do produto ou a apresentação da Cloud TPU (em inglês).

O diagrama a seguir fornece uma visão geral da configuração do cluster distribuído. Em uma única configuração de dispositivo, uma única VM (também chamada de worker de cliente ou CW) alimenta um único worker de TPU (também chamado de worker de serviço ou SW).

No treinamento distribuído, um cluster de CWs e uma fração de pod de TPU correspondente (cluster de SWs) alimentam um único SW. O pipeline de entrada é executado no CW e todo o treinamento do modelo acontece no SW.

image

Objetivos

  • Configurar um grupo de instâncias do Compute Engine e um pod da Cloud TPU para treinamento com PyTorch/XLA
  • Executar treinamento PyTorch/XLA em um pod da Cloud TPU

Antes de começar

Antes de iniciar o treinamento distribuído nos pods da Cloud TPU, verifique se o modelo treina em um único dispositivo da Cloud TPU v2-8 ou v3-8. Caso o modelo tenha problemas significativos de desempenho em um único dispositivo, consulte os guias de práticas recomendadas e de solução de problemas (links em inglês).

Depois de treinar o único dispositivo TPU, execute os passos a seguir para configurar e treinar em um pod da Cloud TPU:

  1. [Opcional] Capture uma imagem de disco da VM em uma imagem de VM.

  2. Crie um modelo de instância a partir de uma imagem de VM.

  3. Crie um grupo de instâncias a partir do seu modelo de instância.

  4. SSH na VM do Compute Engine

  5. Verifique as regras de firewall para permitir a comunicação entre VMs..

  6. Crie um pod da Cloud TPU.

  7. Execute treinamento distribuído no pod.

  8. Faça a limpeza.

[Opcional] Capture uma imagem de disco da VM em uma imagem de VM

É possível utilizar a imagem de disco da VM usada para treinar a única TPU, que já tem o conjunto de dados, pacotes instalados etc., e criar uma imagem de VM (em inglês) que pode ser usada para treinamento de pod. É possível pular esta etapa se não forem necessários pacotes ou outros conjuntos de dados para treinamento e é possível usar uma imagem PyTorch/XLA padrão.

Crie um modelo de instância a partir de uma imagem de VM

Criar um modelo de instância padrão. Ao criar um modelo de instância, use a imagem de VM criada na etapa acima OU use a imagem pública PyTorch/XLA fornecida pelo Google seguindo estas instruções:

  1. Na página Console do Google Cloud, selecione Compute Engine → Modelos de instância.
  2. Na página modelos de instância, clique em Criar modelo de instância.
  3. Na página Criar um modelo de instância, clique no botão Alterar na seção Disco de inicialização.
  4. Se você capturou uma imagem de disco da VM, clique na guia Imagens personalizadas e selecione a imagem capturada. Se você não capturar uma imagem de disco de VM, selecione Deep Learning on Linux no menu suspenso Sistema operacional. Selecione Debian GNU/Linux 9 Stretch + PyTorch/XLA no menu suspenso Versão.
  5. Certifique-se de fazer o seguinte:
    1. Em "Tipo de máquina", selecione n1-standard-16 para este exemplo que usa o treinamento do ResNet-50. Para seu próprio modelo, escolha o tamanho da VM usado para treinar em uma v3-8/v2-8.
    2. Em Acesso à identidade e à API → Acesso aos escopos, selecione Permitir acesso total a todas as APIs Cloud.
  6. Clique em Criar na parte inferior para criar o modelo de instância.

Crie um grupo de instâncias a partir do modelo de instância

  1. No Console do Google Cloud, selecione Grupos de instâncias e clique em Criar grupo de instâncias.
  2. Em Nome, especifique um nome para seu Grupo de instâncias.
  3. Na seção Local, selecione Zona única.
  4. Nas seções Região e Zona, selecione a mesma zona em que o pod da TPU será criado.
  5. Em Modelo de instância, selecione o modelo que você criou na etapa anterior.
  6. Em Número de instâncias: serão necessárias N instâncias em que N equivale ao número total de núcleos do Cloud TPU que você está usando dividido por oito. Neste exemplo, você está usando um pod v2-32, por isso 32 núcleos divididos por 8 fornecem 4. Você precisa de quatro instâncias.
  7. Em Modo de escalonamento automático, selecione Desativado. Se você vir uma lista suspensa Modo de escalonamento automático, selecione Não fazer escalonamento automático.
  8. Em Verificação de integridade, selecione Nenhuma verificação de integridade.
  9. Clique em Criar na parte inferior para criar o grupo de instâncias.

SSH na VM do Compute Engine

Depois de criar o grupo de instâncias, use o SSH na VM do Compute Engine para continuar essas instruções.

  1. No Console do Google Cloud, selecione Compute Engine → Grupos de instâncias e clique no grupo de instâncias que você acabou de criar. Isso abrirá uma janela que exibe todas as instâncias no grupo de instâncias.

  2. No final da linha de descrição de uma das instâncias, clique no botão SSH. Isso abrirá uma janela de sessão da VM. Execute cada comando que começa com (vm)$ na janela de sessão da VM.

Verificar regras de firewall para permitir a comunicação entre VMs

Verifique se as VMs do Compute podem se comunicar na porta 8477. Para isso, consulte as regras de firewall para o projeto OU execute o nmap na VM do Compute Engine.

  1. Execute o comando nmap. Substitua instance-ID por um dos códigos de instância do seu grupo de instâncias. Por exemplo, instance-group-1-g0h2.

    (vm)$ nmap -Pn -p 8477 instance-ID
    
       Starting Nmap 7.40 ( https://nmap.org ) at 2019-10-02 21:35 UTC
       Nmap scan report for pytorch-20190923-n4tx.c.jysohntpu.internal (10.164.0.3)
       Host is up (0.00034s latency).
       PORT     STATE  SERVICE
       8477/tcp closed unknown
    

Contanto que o campo ESTADO não exiba filtrado, as regras de firewall estão configuradas corretamente.

Criar um pod da Cloud TPU

  1. Acesse o Console do Google Cloud e selecione Compute Engine > TPUs.
  2. Clique em Criar nó de TPU.
  3. Em Nome, especifique um nome para o pod de TPU.
  4. Em Zona, especifique a zona a ser usada para o Cloud TPU. Verifique se está na mesma zona que o grupo de instâncias.
  5. Em Tipo de TPU, selecione o Tipo de Cloud TPU. O tipo de Cloud TPU selecionado precisa ter oito vezes o número de instâncias no grupo de instâncias.
  6. Abaixo de Versão do software da TPU, selecione a versão estável mais recente, por exemplo pytorch-1.5.
  7. Use a rede padrão.
  8. Defina o intervalo de endereços IP. Por exemplo, 10.240.0.0.
  9. Clique em Criar para criar o pod de TPU.

Executar treinamentos distribuídos no pod

  1. Na janela de sessão da VM, exporte o nome da Cloud TPU e ative o ambiente conda.

    (vm)$ export TPU_NAME=tpu-pod-name
    (vm)$ conda activate torch-xla-1.5
    
  2. Execute o script de treinamento:

    (torch-xla-1.5)$ python -m torch_xla.distributed.xla_dist \
          --tpu=$TPU_NAME \
          --conda-env=torch-xla-1.5 \
          --env XLA_USE_BF16=1 \
          --env ANY_OTHER=ENV_VAR \
          -- \
          python /usr/share/torch-xla-1.5/pytorch/xla/test/test_train_mp_imagenet.py \
          --fake_data
    

Depois de executar o comando acima, haverá saída semelhante a esta. Observe que o --fake_data está sendo usado. O treinamento leva cerca de meia hora em um pod de TPU v3-32.

2019-09-12 18:24:07  [] Command to distribute: "python" "/usr/share/torch-xla-1.5/pytorch/xla/test/test_train_mp_imagenet.py" "--fake_data"
2019-09-12 18:24:07  [] Cluster configuration: {client_workers: [{10.164.0.24, n1-standard-64, europe-west4-a, tutorial-test-pytorch-pods-05lj}, {10.164.0.36, n1-standard-64, europe-west4-a, tutorial-test-pytorch-pods-5v4b}, {10.164.0.28, n1-standard-64, europe-west4-a, tutorial-test-pytorch-pods-8r6z}, {10.164.0.37, n1-standard-64, europe-west4-a, tutorial-test-pytorch-pods-bhzk}], service_workers: [{10.0.0.3, 8470, v3-32, europe-west4-a, pytorch-1.5}, {10.0.0.2, 8470, v3-32, europe-west4-a, pytorch-1.5}, {10.0.0.4, 8470, v3-32, europe-west4-a, pytorch-1.5}, {10.0.0.5, 8470, v3-32, europe-west4-a, pytorch-1.5}]}
…
2019-09-12 18:24:29 10.164.0.24 [0] 2019-09-12 18:24:29.194350: I tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:196] XRT device (LOCAL) CPU:0 -> /job:c_tpu_worker/replica:0/task:0/device:XLA_CPU:0
2019-09-12 18:24:29 10.164.0.24 [0] 2019-09-12 18:24:29.194423: I tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:196] XRT device (REMOTE) CPU:1 -> /job:c_tpu_worker/replica:0/task:1/device:XLA_CPU:0
2019-09-12 18:24:29 10.164.0.24 [0] 2019-09-12 18:24:29.194431: I tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:196] XRT device (REMOTE) CPU:2 -> /job:c_tpu_worker/replica:0/task:2/device:XLA_CPU:0
2019-09-12 18:24:29 10.164.0.24 [0] 2019-09-12 18:24:29.194437: I tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:196] XRT device (REMOTE) CPU:3 -> /job:c_tpu_worker/replica:0/task:3/device:XLA_CPU:0
2019-09-12 18:24:29 10.164.0.24 [0] 2019-09-12 18:24:29.194443: I tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:196] XRT device (LOCAL) TPU:0 -> /job:c_tpu_worker/replica:0/task:0/device:TPU:0
2019-09-12 18:24:29 10.164.0.24 [0] 2019-09-12 18:24:29.194448: I tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:196] XRT device (LOCAL) TPU:1 -> /job:c_tpu_worker/replica:0/task:0/device:TPU:1
2019-09-12 18:24:29 10.164.0.24 [0] 2019-09-12 18:24:29.194454: I tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:196] XRT device (REMOTE) TPU:10 -> /job:c_tpu_worker/replica:0/task:1/device:TPU:2
…
2019-09-12 18:25:14 10.164.0.24 [0] mesh_shape: 4
2019-09-12 18:25:14 10.164.0.24 [0] mesh_shape: 2
2019-09-12 18:25:14 10.164.0.24 [0] num_tasks: 4
2019-09-12 18:25:14 10.164.0.24 [0] num_tpu_devices_per_task: 8
2019-09-12 18:25:14 10.164.0.24 [0] device_coordinates: 2
2019-09-12 18:25:14 10.164.0.24 [0] device_coordinates: 3
2019-09-12 18:25:14 10.164.0.24 [0] device_coordinates: 0
2019-09-12 18:25:14 10.164.0.24 [0] device_coordinates: 2
…
2019-09-12 18:31:36 10.164.0.28 [2] [xla:2](0) Loss=0.00000 Rate=142.92
2019-09-12 18:31:36 10.164.0.28 [2] [xla:1](0) Loss=0.00000 Rate=116.86
2019-09-12 18:31:36 10.164.0.28 [2] [xla:6](0) Loss=0.00000 Rate=114.17
2019-09-12 18:31:36 10.164.0.28 [2] [xla:4](0) Loss=0.00000 Rate=112.40
2019-09-12 18:31:36 10.164.0.28 [2] [xla:5](0) Loss=0.00000 Rate=101.27
2019-09-12 18:31:36 10.164.0.28 [2] [xla:3](0) Loss=0.00000 Rate=97.01
2019-09-12 18:31:36 10.164.0.28 [2] [xla:7](0) Loss=0.00000 Rate=99.72
2019-09-12 18:31:36 10.164.0.28 [2] [xla:8](0) Loss=0.00000 Rate=98.05
2019-09-12 18:31:36 10.164.0.28 [2] [xla:4](20) Loss=0.00000 Rate=314.58
2019-09-12 18:31:36 10.164.0.28 [2] [xla:3](20) Loss=0.00000 Rate=316.00
2019-09-12 18:31:36 10.164.0.28 [2] [xla:7](20) Loss=0.00000 Rate=317.12
2019-09-12 18:31:36 10.164.0.28 [2] [xla:2](20) Loss=0.00000 Rate=314.21
2019-09-12 18:31:36 10.164.0.28 [2] [xla:6](20) Loss=0.00000 Rate=314.27
2019-09-12 18:31:36 10.164.0.28 [2] [xla:1](20) Loss=0.00000 Rate=311.75
2019-09-12 18:31:36 10.164.0.28 [2] [xla:5](20) Loss=0.00000 Rate=314.76
2019-09-12 18:31:36 10.164.0.28 [2] [xla:8](20) Loss=0.00000 Rate=316.58
2019-09-12 18:31:36 10.164.0.28 [2] [xla:6](40) Loss=0.00000 Rate=423.79
2019-09-12 18:31:36 10.164.0.28 [2] [xla:8](40) Loss=0.00000 Rate=425.21
2019-09-12 18:31:36 10.164.0.28 [2] [xla:5](40) Loss=0.00000 Rate=424.41
2019-09-12 18:31:36 10.164.0.28 [2] [xla:3](40) Loss=0.00000 Rate=423.38
2019-09-12 18:31:36 10.164.0.28 [2] [xla:2](40) Loss=0.00000 Rate=423.19
2019-09-12 18:31:36 10.164.0.28 [2] [xla:7](40) Loss=0.00000 Rate=424.01
2019-09-12 18:31:36 10.164.0.28 [2] [xla:4](40) Loss=0.00000 Rate=422.16
2019-09-12 18:31:36 10.164.0.28 [2] [xla:1](40) Loss=0.00000 Rate=422.09
2019-09-12 18:31:36 10.164.0.28 [2] [xla:1](60) Loss=0.00000 Rate=472.79
2019-09-12 18:31:36 10.164.0.28 [2] [xla:3](60) Loss=0.00000 Rate=472.50
2019-09-12 18:31:36 10.164.0.28 [2] [xla:6](60) Loss=0.00000 Rate=471.26
2019-09-12 18:31:36 10.164.0.28 [2] [xla:8](60) Loss=0.00000 Rate=472.04
2019-09-12 18:31:36 10.164.0.28 [2] [xla:4](60) Loss=0.00000 Rate=471.01
2019-09-12 18:31:36 10.164.0.28 [2] [xla:5](60) Loss=0.00000 Rate=471.77
2019-09-12 18:31:36 10.164.0.28 [2] [xla:2](60) Loss=0.00000 Rate=471.16
2019-09-12 18:31:36 10.164.0.28 [2] [xla:7](60) Loss=0.00000 Rate=471.25
2019-09-12 18:31:36 10.164.0.28 [2] [xla:1](80) Loss=0.00000 Rate=496.79
2019-09-12 18:31:36 10.164.0.28 [2] [xla:4](80) Loss=0.00000 Rate=496.45
2019-09-12 18:31:36 10.164.0.28 [2] [xla:6](80) Loss=0.00000 Rate=496.14
2019-09-12 18:31:36 10.164.0.28 [2] [xla:2](80) Loss=0.00000 Rate=496.45
2019-09-12 18:31:36 10.164.0.28 [2] [xla:3](80) Loss=0.00000 Rate=495.66
2019-09-12 18:31:36 10.164.0.28 [2] [xla:8](80) Loss=0.00000 Rate=496.43
2019-09-12 18:31:36 10.164.0.28 [2] [xla:5](80) Loss=0.00000 Rate=496.22
2019-09-12 18:31:36 10.164.0.28 [2] [xla:7](80) Loss=0.00000 Rate=496.37
2019-09-12 18:31:36 10.164.0.28 [2] [xla:5](100) Loss=0.00000 Rate=503.71
2019-09-12 18:31:36 10.164.0.28 [2] [xla:2](100) Loss=0.00000 Rate=503.13
2019-09-12 18:31:36 10.164.0.28 [2] [xla:3](100) Loss=0.00000 Rate=502.50

Limpar

Para evitar cobranças dos recursos usados neste tutorial na conta do Google Cloud Platform:

Saia da VM do Compute Engine e exclua:

  1. O grupo de instâncias criado em Compute Engine > Grupos de instância

  2. O pod do TPU em Compute Engine > TPUs.

A seguir

Teste as colabs do PyTorch: