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

Follow these steps to authorize the Cloud TPU service account name associated with your Google Cloud project:

  1. Get your Cloud TPU service account name by calling projects.getConfig. Example:

    PROJECT_ID=PROJECT_ID

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

  2. Save the value of the serviceAccountProject and tpuServiceAccount field returned by the API.

  3. Initialize the Cloud TPU service account:

    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

Now add the Cloud TPU service account as a member in your project, with the role Cloud ML Service Agent. Complete the following steps in the Google Cloud console or using the gcloud command:

Console

  1. Log in to the Google Cloud console and choose the project in which you're using the TPU.
  2. Choose IAM & Admin > IAM.
  3. Click the Add button to add a member to the project.
  4. Enter the TPU service account in the Members text box.
  5. Click the Roles dropdown list.
  6. Enable the Cloud ML Service Agent role (Service Agents > Cloud ML Service Agent).

gcloud

  1. Set environment variables containing your project ID and the Cloud TPU service account:

    PROJECT_ID=PROJECT_ID
SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com

  2. Grant the ml.serviceAgent role to the Cloud TPU service account:

    gcloud projects add-iam-policy-binding $PROJECT_ID \
    --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent

For more details about granting roles to service accounts, see the IAM documentation.

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:

TPU v2

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

TPU v3 (Beta)

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:

Pods da TPU v2

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

Pods da TPU v3

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:

Tipo de pod da TPU Número de núcleos de pod disponíveis para uso
TPU_V2_POD 32, 128, 256, 512
TPU_V3_POD 32, 128, 256

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:

URI da imagem do contêiner tpuTfVersion
gcr.io/cloud-ml-public/training/pytorch-xla.1-11 pytorch-1.11
gcr.io/cloud-ml-public/training/pytorch-xla.1-10 pytorch-1.10
gcr.io/cloud-ml-public/training/pytorch-xla.1-9 pytorch-1.9
gcr.io/cloud-ml-public/training/pytorch-xla.1-7 pytorch-1.7
gcr.io/cloud-ml-public/training/pytorch-xla.1-6 pytorch-1.6

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:

TPU v2

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

TPU v3 (Beta)

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

Anterior
Como usar GPUs