Como usar TPUs para treinar seu modelo

As Unidades de Processamento de Tensor (TPUs, na sigla em inglês) são circuitos integrados específicos para aplicativos (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de machine learning. É possível executar jobs de treinamento no AI Platform Training com uma Cloud TPU. O AI Platform Training fornece uma interface de gerenciamento de jobs para que você não precise gerenciar a TPU por conta própria. Em vez disso, use a API jobs do AI Platform Training do mesmo modo que você a usaria para realizar o treinamento em uma CPU ou GPU.

Com as APIs de alto nível do TensorFlow, seus modelos poderão ser executados no hardware da Cloud TPU.

Como configurar o ambiente do Google Cloud

Para configurar o ambiente do Google Cloud, siga as instruções na seção de configuração do guia de primeiros passos.

Como autorizar o Cloud TPU a acessar o projeto

Siga estas etapas para autorizar o nome da conta de serviço do Cloud TPU associada ao seu projeto do Google Cloud:

1. Chame projects.getConfig para receber o nome da conta de serviço do Cloud TPU. Exemplo:

   PROJECT_ID=PROJECT_ID
   
   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
       https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
   

2. Salve os valores dos campos serviceAccountProject e tpuServiceAccount retornados pela API.

3. Inicialize a conta de serviço do Cloud TPU:

   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
     -H "Content-Type: application/json" -d '{}'  \
     https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
   

Agora adicione a conta de serviço da Cloud TPU como membro do seu projeto, com o papel Agente de serviço do Cloud ML. Conclua as etapas a seguir no Console do Google Cloud ou usando o comando gcloud:

Para mais detalhes sobre a concessão de papéis a contas de serviço, consulte a documentação do IAM.

Exemplo: como treinar um modelo MNIST de amostra

Nesta seção, mostramos como treinar um modelo MNIST de amostra usando uma TPU e a versão 2.11 do ambiente de execução. No exemplo de job, o nível de escalonamento predefinido BASIC_TPU é usado para a configuração da máquina. As seções posteriores do guia contêm informações sobre como definir uma configuração personalizada.

Neste exemplo, pressupomos que você esteja usando um shell do Bash com a CLI gcloud instalada. Execute os seguintes comandos para receber o código e enviar o job de treinamento para o AI Platform Training:

1. Faça o download do código dos modelos de referência do TensorFlow e acesse o diretório com o código de amostra:

   git clone https://github.com/tensorflow/models.git \
     --branch=v2.11.0 \
     --depth=1
   
   cd models
   

2. Crie um arquivo setup.py no diretório models. Isso garante que o comando gcloud ai-platform jobs submit training inclua todos os subpacotes necessários no diretório models/official ao criar um tarball do código de treinamento, bem como que o AI Platform Training instale os Conjuntos de dados do TensorFlow como dependências ao executar o job de treinamento. Esse código de treinamento depende dos conjuntos de dados do TensorFlow para carregar os dados do MNIST.

   Para criar o arquivo setup.py, execute o seguinte comando no shell:

   cat << END > setup.py
   from setuptools import find_packages
   from setuptools import setup
   
   setup(
       name='official',
       install_requires=[
          'tensorflow-datasets~=3.1',
          'tensorflow-model-optimization>=0.4.1'
      ],
       packages=find_packages()
   )
   END
   

3. Envie o job de treinamento usando o comando gcloud ai-platform jobs submit training:

   gcloud ai-platform jobs submit training tpu_mnist_1 \
     --staging-bucket=gs://BUCKET_NAME \
     --package-path=official \
     --module-name=official.vision.image_classification.mnist_main \
     --runtime-version=2.11 \
     --python-version=3.7 \
     --scale-tier=BASIC_TPU \
     --region=us-central1 \
     -- \
     --distribution_strategy=tpu \
     --data_dir=gs://tfds-data/datasets \
     --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
   

   Substitua BUCKET_NAME pelo nome de um bucket do Cloud Storage no projeto do Google Cloud. A CLI gcloud envia o código de treinamento empacotado para esse bucket, e o AI Platform Training salva a saída do treinamento no intervalo.

4. Monitore o job de treinamento. Quando o job for concluído, você poderá ver a saída dele no diretório gs://BUCKET_NAME/tpu_mnist_1_output.

Mais detalhes sobre como treinar um modelo na Cloud TPU

Nesta seção, você verá mais informações sobre como configurar um job e treinar um modelo no AI Platform Training com uma Cloud TPU.

Como especificar uma região que ofereça TPUs

Você precisa executar seu job em uma região que tenha TPUs disponíveis. Atualmente, as seguintes regiões dão acesso a TPUs:

• us-central1
• europe-west4

Para conhecer todas as regiões disponíveis aos serviços do AI Platform Training, incluindo treinamento de modelos e previsão on-line/em lote, leia este guia.

Controle de versões do AI Platform Training e do TensorFlow

As versões de ambiente de execução 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9 e 2.11 do AI Platform Training estão disponíveis para treinar seus modelos no Cloud TPU. Saiba mais sobre as versões de ambiente de execução do AI Platform Training e as versões correspondentes do TensorFlow.

A política de controle de versões é a mesma da Cloud TPU. Na sua solicitação de job de treinamento, especifique uma versão de ambiente de execução que esteja disponível para TPUs e corresponda à versão do TensorFlow usada no seu código de treinamento.

Como se conectar ao servidor gRPC da TPU

No programa TensorFlow, use TPUClusterResolver para se conectar ao servidor gRPC da TPU em execução na VM da TPU.

O guia do TensorFlow para usar TPUs mostra como usar TPUClusterResolver com a estratégia de distribuição TPUStrategy.

No entanto, é preciso fazer uma alteração importante ao usar TPUClusterResolver para código executado no AI Platform Training: não forneça nenhum argumento ao criar a instância TPUClusterResolver. Quando os argumentos de palavra-chave tpu, zone e project são definidos com o valor padrão de None, o AI Platform Training fornece automaticamente ao resolvedor de cluster os detalhes de conexão necessários por meio de variáveis de ambiente.

O exemplo a seguir do TensorFlow 2 mostra como inicializar um resolvedor de clusters e uma estratégia de distribuição para treinamento no AI Platform Training:

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

Como usar TPUs no código do TensorFlow

Para usar as TPUs em uma máquina, use a API TPUStrategy do TensorFlow 2. O guia do TensorFlow para usar TPUs mostra como fazer isso.

Para treinar com TPUs no TensorFlow 1, use a API TPUEstimator. O guia do Cloud TPU para a API TPUEstimator mostra como fazer isso.

A documentação do Cloud TPU também fornece uma lista de operações de nível baixo do TensorFlow disponíveis no Cloud TPU.

Como usar TPUs no código do PyTorch

Para utilizar uma TPU ao aplicar um contêiner de PyTorch pré-construído, use o pacote torch_xla. Saiba como usar torch_xla para TPU no treinamento na documentação do PyTorch. Para mais exemplos de uso de torch_xla, consulte os tutoriais no repositório GitHub do PyTorch do GitHub.

Ao treinar usando uma TPU no AI Platform Training, você usa um único dispositivo XLA, e não vários dispositivos XLA.

Consulte também a seção a seguir nesta página sobre como configurar o job de treinamento para PyTorch e TPU.

Como configurar uma máquina personalizada de TPU

Um job de treinamento de TPU é executado em uma configuração de duas VMs. Uma VM (mestre) executa seu código Python. O mestre controla o servidor do TensorFlow em execução em um worker de TPU.

Para usar uma TPU com o AI Platform Training, configure seu job de treinamento para acessar uma máquina ativada para TPU de uma destas três maneiras:

• Use o nível de escalonamento BASIC_TPU. Use esse método para acessar os aceleradores da TPU v2.
• Use um worker cloud_tpu e um tipo de máquina legada para a VM principal. Use esse método para acessar os aceleradores da TPU v2.
• Use um worker cloud_tpu e um tipo de máquina do Compute Engine para a VM principal. É possível usar esse método para acessar os aceleradores TPU v2 ou TPU v3. Os aceleradores de TPU v3 estão disponíveis na versão Beta.

Máquina ativada para TPU básica

Defina o nível de escalonamento como BASIC_TPU para conseguir uma VM principal e uma VM de TPU, incluindo uma TPU com oito núcleos TPU v2, como você fez ao executar o exemplo anterior.

Worker de TPU em uma configuração de tipo de máquina legada

Como alternativa, você pode definir uma configuração de máquina personalizada, caso precise de mais recursos de computação na VM mestre:

• Defina o nível de escalonamento como CUSTOM.
• Configure a VM principal para usar um tipo de máquina legada adequado aos requisitos do job.
• Configure workerType como cloud_tpu para conseguir uma VM de TPU, incluindo um Cloud TPU com oito núcleos TPU v2.
• Defina workerCount como 1.
• Não especifique um servidor de parâmetros ao usar um Cloud TPU. O serviço rejeitará a solicitação de job se parameterServerCount for maior que zero.

No exemplo a seguir, mostramos um arquivo config.yaml que usa esse tipo de configuração:

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

Worker TPU em uma configuração de tipo de máquina do Compute Engine

Também é possível definir uma configuração de máquina personalizada com um tipo de máquina do Compute Engine para sua VM mestre e um acceleratorConfig anexado à sua VM de TPU.

É possível usar esse tipo de configuração para configurar um worker da TPU com oito núcleos TPU v2 (semelhante a uma configuração sem acceleratorConfig) ou um worker da TPU com oito núcleos TPU v3 (Beta). Leia mais sobre a diferença entre os aceleradores TPU v2 e TPU v3.

Usar um tipo de máquina do Compute Engine também oferece mais flexibilidade para configurar sua VM mestre:

• Defina o nível de escalonamento como CUSTOM.
• Configure a VM mestre para usar um tipo de máquina do Compute Engine que atenda aos seus requisitos de trabalho.
• Defina workerType como cloud_tpu.
• Adicione um workerConfig com um campo acceleratorConfig. Dentro de acceleratorConfig, defina type como TPU_V2 ou TPU_V3 e count como 8. Não anexe qualquer outro número de núcleos de TPU.
• Defina workerCount como 1.
• Não especifique um servidor de parâmetros ao usar um Cloud TPU. O serviço rejeitará a solicitação de job se parameterServerCount for maior que zero.

No exemplo a seguir, mostramos um arquivo config.yaml que usa esse tipo de configuração:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3
      count: 8

Como usar pods de TPU

Um Pod de TPU é um conjunto de dispositivos de TPU conectados por interfaces de rede dedicadas de alta velocidade. Um Pod de TPU pode ter até 2.048 núcleos, permitindo distribuir a carga de processamento entre várias TPUs.

Para usar pods de TPU, primeiro é necessário registrar uma solicitação de aumento de cota.

Os seguintes exemplos de arquivos config.yaml mostram como usar pods de TPU:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2_POD
      count: 128

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3_POD
      count: 32

Há limitações no número de núcleos de pod que podem ser usados para cada tipo de TPU. Configurações disponíveis:

Para mais detalhes sobre como fazer uso total dos núcleos de TPU, consulte a documentação do Cloud TPU sobre pods de TPU.

Como usar um contêiner do PyTorch criado anteriormente em um worker da TPU

Se você quiser realizar o treinamento PyTorch com uma TPU, especifique o campo tpuTfVersion no trainingInput do job de treinamento. Defina tpuTfVersion para corresponder à versão do contêiner PyTorch pré-criado que você está usando para o treinamento.

O AI Platform Training é compatível com treinamento de TPUs para os seguintes contêineres pré-criados do PyTorch:

Por exemplo, para treinar usando o contêiner pré-criado PyTorch 1.11, use o seguinte arquivo config.yaml para configurar o treinamento:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
    tpuTfVersion: pytorch-1.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

Consulte também a seção anterior nesta página sobre Como usar TPUs no código PyTorch.

Como usar um contêiner personalizado em um worker de TPU

Se você quiser executar um contêiner personalizado no seu worker de TPU, em vez de usar uma das versões de ambiente de execução do AI Platform Training compatíveis com TPUs, especifique outro campo de configuração quando enviar seu job de treinamento. Defina tpuTfVersion como uma versão do ambiente de execução que inclua a versão do TensorFlow que o contêiner usa. Especifique uma versão do ambiente de execução compatível para treinamento com TPUs.

Como você está configurando o job para usar um contêiner personalizado, o AI Platform Training não usa o ambiente de execução ao executar o job de treinamento. No entanto, ele precisa que esse campo seja preenchido para preparar adequadamente o worker da TPU para a versão do TensorFlow usada pelo contêiner personalizado.

No exemplo a seguir, mostramos um arquivo config.yaml com uma configuração de TPU semelhante àquela da seção anterior, exceto que neste caso a VM principal e o worker da TPU executam contêineres personalizados diferentes:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V3
      count: 8

Se você usar o comando gcloud beta ai-platform jobs submit training para enviar o job de treinamento, poderá especificar o campo da API tpuTfVersion com a sinalização --tpu-tf-version, em vez de em um arquivo config.yaml.

Como usar TPUClusterResolver após o provisionamento da TPU

Ao usar um contêiner personalizado, você precisa aguardar o provisionamento da TPU antes de chamar TPUClusterResolver. O código de amostra a seguir mostra como lidar com a lógica TPUClusterResolver:

def wait_for_tpu_cluster_resolver_ready():
  """Waits for `TPUClusterResolver` to be ready and return it.

  Returns:
    A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
    return None.
  Raises:
    RuntimeError: if failed to schedule TPU.
  """
  tpu_config_env = os.environ.get('TPU_CONFIG')
  if not tpu_config_env:
    tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
    return None

  tpu_node = json.loads(tpu_config_env)
  tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)

  num_retries = 40
  for i in range(num_retries):
    try:
      tpu_cluster_resolver = (
          tf.contrib.cluster_resolver.TPUClusterResolver(
              tpu=[tpu_node['tpu_node_name']],
              zone=tpu_node['zone'],
              project=tpu_node['project'],
              job_name='worker'))
      tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
      if 'worker' in tpu_cluster_resolver_dict:
        tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
        return tpu_cluster_resolver
    except Exception as e:
      if i < num_retries - 1:
        tf.logging.info('Still waiting for provisioning of TPU VM instance.')
      else:
        # Preserves the traceback.
        raise RuntimeError('Failed to schedule TPU: {}'.format(e))
    time.sleep(10)

  # Raise error when failed to get TPUClusterResolver after retry.
  raise RuntimeError('Failed to schedule TPU.')

Saiba mais sobre o treinamento distribuído com contêineres personalizados.

A seguir

• Saiba mais sobre modelos de treinamento no AI Platform Training.
• Saiba mais sobre como ajustar hiperparâmetros no AI Platform Training, com atenção especial aos detalhes de ajuste com a Cloud TPU.
• Veja outros modelos de referência do Cloud TPU.
• Otimize os modelos para a Cloud TPU seguindo as práticas recomendadas para Cloud TPU.
• Para receber ajuda sobre diagnósticos e correções, consulte a solução de problemas e as perguntas frequentes do Cloud TPU.