Como treinar modelos PyTorch em pods da 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 e mostra como o treinamento acontece. Na configuração de um único dispositivo, uma única VM (worker de cliente = CW, na sigla em inglês) alimenta um único worker de TPU (worker de serviço = SW, na sigla em inglês). Da mesma forma, no treinamento distribuído, um cluster de VMs/CWs e também uma fração de pod do TPU correspondente (cluster de SWs) e cada um dos CWs alimenta um único SW. O pipeline de entrada é executado no CW e todo o treinamento do modelo acontece no SW.

imagem

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. Pule esta etapa se não precisar de mais pacotes ou conjuntos de dados para treinamento e puder usar apenas 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 no passo anterior OU a imagem pública PyTorch/XLA fornecida pelo Google. Para fazer isso, siga estas instruções:

  1. Na página "Criar um modelo de instância", clique no botão “Alterar” na seção “Disco de inicialização”.
  2. Se você capturou uma imagem de disco da VM, clique na guia "Imagens personalizadas" e selecione a imagem capturada. Se você não capturou uma imagem de disco da VM, selecione a imagem pública PyTorch/XLA no menu suspenso "Imagens do SO".
  3. 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 do Cloud".
  4. Clique em "Selecionar" na parte inferior para criar o modelo da instância.

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

No Console do Google Cloud Platform, selecione Compute Engine > Grupos de instâncias para acessar o formulário Criar um grupo de instâncias. Ao preencher o formulário, especifique os seguintes parâmetros para a configuração:

  1. Em "Nome", especifique um nome para o grupo de instâncias.
  2. Na seção "Local", selecione "Zona única".
  3. Nas seções "Região" e "Zona", selecione a mesma zona em que o pod do TPU será criado.
  4. Em "Modelo de instância", selecione o modelo criado na etapa anterior.
  5. Em "Modo de escalonamento automático", selecione "Desativado". Se você vir o menu suspenso "Modo de escalonamento automático", selecione "Não configurar o escalonamento automático".
  6. Em "Número de instâncias", será preciso N instâncias em que N é igual ao número total de núcleos usados dividido por 8, já que cada instância de VM (CW) alimenta oito núcleos de TPU. Neste exemplo, N é igual a 4, então você precisará de 4 instâncias.
  7. Em "Verificação de integridade", selecione "Sem verificação de integridade".
  8. Clique em "Selecionar" 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 Platform, selecione Compute Engine > Grupos de instâncias e clique no grupo de instâncias recém criado. 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. A variável instance-ID é uma das instâncias 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

Acesse Google Cloud Platform Console e selecione Compute Engine > TPUs. Isso exibirá a página em que é possível criar um nó de TPU e especificar os seguintes parâmetros:

  1. Em "Nome", especifique um nome para o pod de TPU.
  2. Em “Zona”, especifique a zona a ser usada para a Cloud TPU. Verifique se está na mesma zona que o grupo de instâncias.
  3. Em "Tipo de TPU", selecione o tipo de Cloud TPU.
  4. Em "Versão do software de TPU", selecione a versão estável mais recente (pytorch-0.X), por exemplo, pytorch-0.5.
  5. Usar a rede padrão.
  6. Defina o intervalo de endereços IP. Por exemplo, 10.240.0.0.

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-0.5
    
  2. Execute o script de treinamento:

    (torch-xla-0.5)$ python -m torch_xla.distributed.xla_dist \
          --tpu=$TPU_NAME \
          --conda-env=torch-xla-0.5 \
          --env XLA_USE_BF16=1 \
          --env ANY_OTHER=ENV_VAR \
          -- python /usr/share/torch-xla-0.5/pytorch/xla/test/test_train_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-0.5/pytorch/xla/test/test_train_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-0.5}, {10.0.0.2, 8470, v3-32, europe-west4-a, pytorch-0.5}, {10.0.0.4, 8470, v3-32, europe-west4-a, pytorch-0.5}, {10.0.0.5, 8470, v3-32, europe-west4-a, pytorch-0.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

Limpeza

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

Use as colabs do PyTorch:

Esta página foi útil? Conte sua opinião sobre:

Enviar comentários sobre…

Precisa de ajuda? Acesse nossa página de suporte.