Introdução ao Cloud TPU

As Tensor Processing Units (TPUs) são circuitos integrados específicos da aplicação (ASICs) desenvolvidos pela Google e usados para acelerar as cargas de trabalho de aprendizagem automática. Para mais informações sobre o hardware da TPU, consulte o artigo Arquitetura da TPU. O Cloud TPU é um serviço Web que disponibiliza TPUs como recursos de computação escaláveis no Google Cloud.

As UTPs preparam os seus modelos de forma mais eficiente através de hardware concebido para realizar operações de matrizes grandes, que se encontram frequentemente em algoritmos de aprendizagem automática. As TPUs têm memória de alta largura de banda (HBM) no chip, o que lhe permite usar modelos maiores e tamanhos de lotes maiores. As TPUs podem ser ligadas em grupos denominados fatias que escalam as suas cargas de trabalho com poucas ou nenhuma alteração ao código.

O código executado em TPUs tem de ser compilado pelo compilador de álgebra linear do acelerador (XLA). O XLA é um compilador just-in-time que usa o gráfico emitido por uma aplicação de framework de ML e compila os componentes de álgebra linear, perda e gradiente do gráfico em código de máquina da TPU. O resto do programa é executado na máquina anfitriã da TPU. O compilador XLA faz parte da imagem da VM da TPU que é executada numa máquina anfitriã da TPU.

Para mais informações sobre as unidades de processamento tensor, consulte o artigo Como pensar nas TPUs.

Quando usar TPUs

As Cloud TPUs estão otimizadas para cargas de trabalho específicas. Em algumas situações, pode querer usar GPUs ou CPUs em instâncias do Compute Engine para executar as suas cargas de trabalho de aprendizagem automática. Em geral, pode decidir que hardware é mais adequado para a sua carga de trabalho com base nas diretrizes que se seguem.

CPUs

  • Criação rápida de protótipos que requer a máxima flexibilidade
  • Modelos simples que não demoram muito tempo a preparar
  • Modelos pequenos com tamanhos de lotes pequenos e eficazes
  • Modelos que contêm muitas operações personalizadas do TensorFlow escritas em C++
  • Modelos limitados pela E/S disponível ou pela largura de banda de rede do sistema anfitrião

GPUs

  • Modelos com um número significativo de operações personalizadas do PyTorch/JAX que têm de ser executadas, pelo menos, parcialmente em CPUs
  • Modelos com operações do TensorFlow que não estão disponíveis na Cloud TPU (consulte a lista de operações do TensorFlow disponíveis)
  • Modelos médios a grandes com tamanhos de lotes efetivos maiores

TPUs

  • Modelos dominados por cálculos de matrizes
  • Modelos sem operações personalizadas do PyTorch/JAX no ciclo de preparação principal
  • Modelos que são preparados durante semanas ou meses
  • Modelos grandes com tamanhos de lotes efetivos grandes
  • Modelos com incorporações ultragrandes comuns em classificações avançadas e cargas de trabalho de recomendações

As Cloud TPUs não são adequadas para as seguintes cargas de trabalho:

  • Programas de álgebra linear que requerem ramificações frequentes ou contêm muitas operações de álgebra ao nível do elemento
  • Cargas de trabalho que requerem aritmética de alta precisão
  • Cargas de trabalho de redes neurais que contêm operações personalizadas no ciclo de preparação principal

TPUs no Google Cloud

Pode usar TPUs através de VMs do Cloud TPU, do Google Kubernetes Engine e do Vertex AI. A tabela seguinte apresenta os recursos de cada Google Cloud serviço.

Google Cloud serviço Recursos
Cloud TPU Comece a usar as VMs da Cloud TPU
Google Kubernetes Engine

Acerca das TPUs no GKE

Execute o Ray no GKE com TPUs

Vertex AI

Preparação no Vertex AI com TPUs

Use TPUs para a previsão online no Vertex AI

Práticas recomendadas para o desenvolvimento de modelos

Um programa cuja computação seja dominada por operações não matriciais, como add, reshape ou concatenate, provavelmente não vai alcançar uma elevada utilização da MXU. Seguem-se algumas diretrizes para ajudar a escolher e criar modelos adequados para a Cloud TPU.

Esquema

O compilador XLA realiza transformações de código, incluindo a divisão de uma multiplicação de matrizes em blocos mais pequenos, para executar eficientemente os cálculos na unidade de matrizes (MXU). A estrutura do hardware MXU, uma matriz sistólica de 128 x 128, e o design do subsistema de memória das TPUs, que prefere dimensões que sejam múltiplos de 8, são usados pelo compilador XLA para a eficiência da segmentação.

Consequentemente, determinados esquemas são mais propícios à disposição em mosaico, enquanto outros requerem remodelações antes de poderem ser dispostos em mosaico. As operações de reformulação têm frequentemente limites de memória na Cloud TPU.

Formas

O compilador XLA compila um gráfico de ML a tempo do primeiro lote. Se quaisquer lotes subsequentes tiverem formas diferentes, o modelo não funciona. (A recompilação do gráfico sempre que a forma muda é demasiado lenta.) Por conseguinte, qualquer modelo que tenha tensores com formas dinâmicas não é adequado para as TPUs.

Preenchimento

Um programa de TPU na nuvem de elevado desempenho é aquele em que a computação densa pode ser dividida em blocos de 128 x 128. Quando um cálculo de matriz não consegue ocupar uma MXU inteira, o compilador preenche os tensores com zeros. Existem duas desvantagens no preenchimento:

  • Os tensores preenchidos com zeros não utilizam totalmente o núcleo da TPU.
  • O preenchimento aumenta a quantidade de armazenamento de memória no chip necessário para um tensor e pode levar a um erro de falta de memória no caso extremo.

Embora o preenchimento seja realizado automaticamente pelo compilador XLA quando necessário, pode determinar a quantidade de preenchimento realizado através da ferramenta op_profile. Pode evitar o preenchimento escolhendo dimensões de tensores adequadas para TPUs.

Dimensões

A escolha de dimensões de tensores adequadas contribui significativamente para extrair o máximo desempenho do hardware da TPU, particularmente da MXU. O compilador XLA tenta usar o tamanho do lote ou uma dimensão de funcionalidade para usar ao máximo a MXU. Por conseguinte, um destes valores tem de ser um múltiplo de 128. Caso contrário, o compilador preenche um deles até 128. Idealmente, o tamanho do lote e as dimensões das caraterísticas devem ser múltiplos de 8, o que permite extrair um elevado desempenho do subsistema de memória.

Introdução às Cloud TPUs

  1. Configure uma Google Cloud conta
  2. Ative a API Cloud TPU
  3. Conceda acesso do Cloud TPU aos seus contentores do Cloud Storage
  4. Execute um cálculo básico numa TPU
  5. Prepare um modelo de referência numa TPU
  6. Analise o seu modelo

Pedir ajuda

Para receber ajuda, contacte o apoio técnico do Cloud TPU. Se tiver um Google Cloud projeto ativo, prepare-se para facultar as seguintes informações:

  • O seu Google Cloud ID do projeto
  • O nome da sua TPU, se existir
  • Outras informações que quer facultar

O que se segue?

Quer saber mais sobre a Cloud TPU? Os seguintes recursos podem ajudar: