Introdução ao Cloud TPU

Unidades de processamento de tensor ou TPUs (na sigla em inglês) são aceleradores de hardware desenvolvidos pelo Google para cargas de trabalho de machine learning. Para saber mais sobre o hardware de TPUs, consulte Arquitetura do sistema. O Cloud TPU é um serviço da Web que disponibiliza TPUs como recursos de computação escalonáveis no Google Cloud.

As TPUs treinam seus modelos com mais eficiência se usarem hardwares projetados para executar operações de matriz de grande escala, geralmente encontradas em algoritmos de machine learning. As TPUs têm memória de largura de banda alta (HBM, na sigla em inglês) que permite usar modelos maiores e tamanhos de lote. As TPUs podem ser conectadas em grupos chamados pods, que escalonam suas cargas de trabalho com pouca ou nenhuma alteração no código.

Recomendamos iniciar seu projeto de machine learning com uma única TPU e escalonar para um Pod de TPU na produção. Para começar, aproveite um conjunto de modelos de referência de código aberto que as equipes de pesquisa e engenharia do Google otimizam para usar com TPUs. Para mais informações, consulte Modelos de referência.

Como funciona?

Para entender como as TPUs funcionam, vale a pena entender como outros aceleradores lidam com os desafios computacionais do treinamento de modelos de ML.

Como funciona uma CPU

Uma CPU é um processador de uso geral com base na arquitetura de von Neumann. Isso significa que uma CPU trabalha com software e memória, assim:

Ilustração de como funciona uma CPU

O maior benefício das CPUs é a flexibilidade. É possível carregar qualquer tipo de software em uma CPU para vários tipos diferentes de aplicativos. Por exemplo, é possível usar uma CPU para processar palavras em um PC, controlar motores de foguetes, executar transações bancárias ou classificar imagens com uma rede neural.

Uma CPU carrega valores da memória, realiza um cálculo com base nos valores e armazena o resultado de cada cálculo na memória. O acesso à memória é lento em comparação com a velocidade de cálculo e pode limitar a capacidade total de CPUs. Isso geralmente é chamado de gargalo de von Neumann.

Como funciona uma GPU

Para aumentar a capacidade de processamento, as GPUs contêm milhares de unidades lógicas aritméticas (ALUs, na sigla em inglês) em um único processador. Uma GPU moderna geralmente contém entre 2.500 e 5.000 ALUs. O grande número de processadores significa que você pode executar milhares de multiplicações e adições simultaneamente.

Ilustração de como funciona uma GPU

Essa arquitetura de GPU funciona bem em aplicativos com enorme paralelismo, como a operação de matrizes em uma rede neural. Na verdade, em uma carga de trabalho de treinamento típica para aprendizado profundo, uma GPU fornece uma ordem de grandeza maior do que uma CPU.

Mas a GPU ainda é um processador de uso geral que precisa ser compatível com muitos aplicativos e softwares diferentes. Portanto, as GPUs têm o mesmo problema que as CPUs. Para cada cálculo nas milhares de ALUs, uma GPU precisa acessar registros ou memória compartilhada para ler e armazenar os resultados de cálculos intermediários.

Como funciona uma TPU

O Google projetou as Cloud TPUs como um processador de matriz especializado em cargas de trabalho de redes neurais. As TPUs não executam processadores de texto, controlam mecanismos de foguete ou executam transações bancárias, mas podem processar operações de matriz maciças usadas em redes neurais em velocidades rápidas.

A principal tarefa das TPUs é o processamento de matrizes, que é uma combinação de operações de multiplicação e acumulação. As TPUs contêm milhares de acumuladores de multiplicação que estão diretamente conectados entre si para formar uma matriz física grande. Isso é chamado de arquitetura de matriz sistólica. O Cloud TPU v3 contém duas matrizes sistólicas de ALUs de 128 x 128, em um único processador.

O host da TPU transmite os dados para uma fila de entrada. A TPU carrega dados da fila de alimentação e os armazena na memória HBM. Quando a computação é concluída, a TPU carrega os resultados na fila de saída. Em seguida, o host da TPU lê os resultados da fila de saída e os armazena na memória do host.

Para executar as operações da matriz, a TPU carrega os parâmetros da memória HBM na MXU.

Ilustração de como uma TPU carrega parâmetros da memória

Em seguida, ela carrega dados da memória HBM. À medida que cada multiplicação é executada, o resultado é transmitido para o próximo multiplicador de acumulador. A saída é a soma de todos os resultados de multiplicação entre os dados e os parâmetros. Não é necessário acessar a memória durante o processo de multiplicação de matrizes.

Ilustração de como uma TPU carrega dados da memória

Como resultado, as TPUs alcançam uma alta capacidade de computação em cálculos de redes neurais.

Quando usar TPUs

Os Cloud TPUs são otimizados para cargas de trabalho específicas. Em algumas situações, é recomendável usar GPUs ou CPUs em instâncias do Compute Engine para executar cargas de trabalho de aprendizado de máquina. Em geral, é possível decidir qual hardware é melhor para sua carga de trabalho com base nas diretrizes a seguir:

CPUs

  • prototipagem rápida que requer flexibilidade máxima
  • modelos simples que não levam muito tempo para serem treinados
  • modelos pequenos com tamanhos de lotes menores
  • Modelos com muitas operações personalizadas do TensorFlow/PyTorch/JAX escritas em C++
  • modelos limitados por E/S disponíveis ou pela largura de banda de rede do sistema host

GPUs

  • modelos com um número significativo de operações personalizadas do TensorFlow/PyTorch/JAX que precisam ser executadas pelo menos parcialmente nas CPUs.
  • modelos com operações do TensorFlow/PyTorch que não estão disponíveis no Cloud TPU
  • modelos de médio a grande porte com tamanhos de lote maiores

TPUs

  • modelos com preponderância de computações matriciais
  • modelos sem operações personalizadas do TensorFlow/PyTorch/JAX no loop de treinamento principal
  • modelos com treinamento de semanas ou meses
  • modelos grandes com tamanhos grandes de lote efetivo

Os Cloud TPUs não são adequados para as seguintes cargas de trabalho:

  • Programas de álgebra linear que exigem ramificações frequentes ou contêm muitas operações de álgebra com elementos
  • Cargas de trabalho que acessam a memória de maneira esparsa
  • Cargas de trabalho que exigem aritmética de alta precisão
  • Cargas de trabalho de rede neural com operações personalizadas no loop de treinamento principal

Integração com o VPC Service Controls

O VPC Service Controls do Cloud TPU permite definir perímetros de segurança em torno dos recursos do Cloud TPU e controlar a movimentação de dados no limite do perímetro. Para saber mais sobre o VPC Service Controls, consulte a visão geral sobre ele. Para saber mais sobre as limitações de uso da Cloud TPU com os VPC Service Controls, consulte os produtos e limitações com suporte.

Primeiros passos com o Cloud TPU

Configure uma conta do Google Cloud Antes de usar os recursos da Cloud TPU, é necessário criar uma conta e um projeto do Google Cloud.
Ative as APIs Cloud TPU Para treinar um modelo, ative as APIs Compute Engine e Cloud TPU.
Conceda ao Cloud TPU acesso a buckets do Cloud Storage A Cloud TPU requer acesso aos buckets do Cloud Storage em que os conjuntos de dados são armazenados.
Treine o modelo Leia um dos guias de início rápido ou tutoriais do Cloud TPU para começar.
Analisar seu modelo Use o TensorBoard ou outras ferramentas para visualizar seu modelo e rastrear as principais métricas.

A seguir

Quer saber mais sobre o Cloud TPU? Os recursos a seguir podem ajudar.

Guia de início rápido usando o Compute Engine Use um de nossos guias de início rápido para tentar treinar um modelo com o Cloud TPU.
Colabs TPU Testar a Cloud TPU usando Colabs gratuitos.
Tutoriais da Cloud TPU Testar o Cloud TPU usando modelos de ML de referência.
Arquitetura do sistema do Cloud TPU Veja informações mais detalhadas sobre TPUs.
Preços Saiba como o Cloud TPU pode processar as cargas de trabalho de machine learning de maneira econômica.
Entre em contato com a equipe de vendas Tem uma implementação ou um aplicativo específico para discutir? Entre em contato com seu representante de vendas da Cloud TPU.