Cloud Tensor Processing Units (Cloud TPUs)

As unidades de processamento de tensor (TPUs, na sigla em inglês) são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de aprendizado de máquina. Todo o desenvolvimento das TPUs foi beneficiado pela extensa experiência e pela liderança do Google em machine learning.

Com o Cloud TPU, é possível executar cargas de trabalho de machine learning no hardware acelerador de TPUs do Google usando o TensorFlow. O Cloud TPU foi desenvolvido para proporcionar o máximo de desempenho e flexibilidade, ajudando pesquisadores, desenvolvedores e empresas a criar clusters de computação do TensorFlow que podem usar CPUs, GPUs e TPUs. As APIs de alto nível do TensorFlow ajudam a executar modelos no hardware do Cloud TPU.

Vantagens das TPUs

Os recursos do Cloud TPU aceleram o desempenho da computação algébrica linear, que é usada intensamente em aplicativos de aprendizado de máquina. As TPUs minimizam o tempo que se leva para gerar precisão ao treinar modelos de rede neural grandes e complexos. Modelos que anteriormente levavam semanas para serem treinados em outras plataformas de hardware podem convergir em horas nas TPUs.

Disponibilidade

As regiões e zonas do Cloud TPU mostram os locais em que os Cloud TPUs estão disponíveis.

Modelo de programação do Cloud TPU

Os Cloud TPUs são muito rápidos para executar computações de vetor e matriz densas. A transferência de dados entre o Cloud TPU e a memória host é lenta em comparação com a velocidade da computação: a velocidade do barramento PCIe é muito mais lenta do que a interconexão do Cloud TPU e a memória de alta largura de banda (HBM, na sigla em inglês) on-chip. A compilação parcial de um modelo, em que a execução alterna entre o host e o dispositivo, faz com que a TPU fique ociosa na maior parte do tempo, aguardando a chegada dos dados pelo barramento PCIe. Para aliviar essa situação, o modelo de programação do Cloud TPU foi projetado para executar grande parte do treinamento na TPU, se possível, todo o loop de treinamento.

Estes são alguns dos principais recursos do modelo de programação da TPU:

  • Todos os parâmetros do modelo são mantidos na memória de alta largura de banda on-chip.
  • O custo de iniciar computações no Cloud TPU é amortizado pela execução de muitas etapas de treinamento em loop.
  • Os dados de treinamento de entrada são transmitidos para uma fila de "infeed" no Cloud TPU. Um programa em execução no Cloud TPU recupera os lotes dessas filas durante cada etapa de treinamento.
  • O servidor do TensorFlow em execução na máquina host (a CPU conectada ao dispositivo do Cloud TPU) busca dados e faz o pré-processamento deles antes de "alimentá-los" no hardware do Cloud TPU.
  • Paralelismo de dados: os núcleos de um Cloud TPU executam de maneira síncrona um programa idêntico que reside nas respectivas HBMs. É realizada uma operação de redução ao final de cada etapa da rede neural em todos os núcleos.

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 preponderância de operações personalizadas do TensorFlow escritas em C++
    • modelos limitados por E/S disponíveis ou pela largura de banda de rede do sistema host

  • GPUs

    • modelos com código-fonte inexistente ou oneroso demais para ser alterado
    • modelos com um número significativo de operações personalizadas do TensorFlow que precisam ser executadas pelo menos parcialmente nas CPUs
    • modelos com operações do TensorFlow que não estão disponíveis no Cloud TPU (consulte a lista de operações do TensorFlow disponíveis)
    • 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 do TensorFlow personalizadas no loop de treinamento principal
    • modelos com treinamento de semanas ou meses
    • modelos grandes e muito grandes com tamanhos de lote imensos

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

  • Programas algébricos lineares que requerem ramificação frequente ou com preponderância de operações algébricas de elemento. As TPUs são otimizadas para realizar multiplicações matriciais pesadas e rápidas. Portanto, é improvável que uma carga de trabalho que não seja composta majoritariamente por multiplicações matriciais tenha um bom desempenho nas TPUs, em comparação a outras plataformas.
  • As cargas de trabalho que acessam a memória de maneira esparsa talvez não estejam disponíveis nas TPUs.
  • Cargas de trabalho que exigem aritmética de alta precisão. Por exemplo, a aritmética de precisão dupla não é adequada para TPUs.
  • Cargas de trabalho de rede neural com operações personalizadas do TensorFlow escritas em C++. Especificamente, as operações personalizadas no corpo do loop de treinamento principal não são adequadas para as TPUs.

As cargas de trabalho de rede neural precisam conseguir executar várias iterações de todo o loop de treinamento na TPU. Embora esse não seja um requisito fundamental das próprias TPUs, essa é uma das restrições atuais do ecossistema de software da TPU que é necessária para a eficiência.

Diferenças em comparação ao treinamento convencional

Um gráfico de treinamento do TensorFlow típico consiste em vários subgráficos sobrepostos que fornecem uma variedade de funcionalidades, incluindo:

  • operações de E/S para ler os dados de treinamento;
  • estágios de pré-processamento de entradas, muitas vezes conectados por meio de filas;
  • variáveis do modelo;
  • código de inicialização para as variáveis;
  • o próprio modelo;
  • funções de perda;
  • código do gradiente (normalmente gerado automaticamente);
  • operações "summary" para o monitoramento do treinamento;
  • operações para salvar/restaurar os pontos de verificação.

No Cloud TPU, os programas do TensorFlow são compilados pelo compilador just-in-time XLA. Ao treinar no Cloud TPU, o único código que pode ser compilado e executado no hardware é aquele que corresponde às partes densas dos subgráficos do modelo, de perda e de gradiente. Todas as outras partes do programa do TensorFlow são executadas nas máquinas host (servidor do Cloud TPU) como parte de uma sessão distribuída normal do TensorFlow. Isso normalmente consiste em operações de E/S que leem dados de treinamento, todos os códigos de pré-processamento (por exemplo, decodificação de imagens compactadas, amostragem/corte em ordem aleatória, montagem de minilotes de treinamento) e todas as partes de manutenção do gráfico, como salvar/restaurar o checkpoint.

Práticas recomendadas para o desenvolvimento de modelos

Um único chip do Cloud TPU tem dois núcleos, cada um com várias unidades matriciais (MXUs) feitas para acelerar programas dominados por convoluções e multiplicações matriciais densas (consulte Arquitetura de sistema). Programas que dedicam uma parcela considerável do tempo de execução à realização de multiplicações matriciais costumam ser adequados para o Cloud TPU. Um programa com computação dominada por operações não matriciais, como "add", "reshape" ou "concatenate", provavelmente não atingirá uma alta utilização de MXU. A seguir, conheça algumas diretrizes que ajudam a escolher e criar modelos adequados para o Cloud TPU.

Um único dispositivo do Cloud TPU consiste em quatro chips, cada um com dois núcleos de TPU. Portanto, para uma utilização eficiente do Cloud TPU, um programa deve fazer uso de cada um dos oito núcleos.

Layout

O compilador XLA executa transformações no código, incluindo a divisão do produto de matrizes em blocos menores para executar computações na MXU com mais eficiência. A estrutura do hardware da MXU (um arranjo sistólico de 128x128) e o design do subsistema de memória da TPU (que preferencialmente usa dimensões que são múltiplas de oito) são usados pelo compilador XLA para aumentar a eficiência da divisão em blocos. Por consequência, determinados layouts são mais propícios para a divisão em blocos, enquanto outros requerem a realização de remodelações antes da divisão. As operações de remodelação geralmente são limitadas pela memória no Cloud TPU.

Formas

O compilador XLA compila um gráfico do TensorFlow a tempo para o primeiro lote. Se os lotes subsequentes tiverem formas diferentes, o modelo não funcionará. Compilar novamente o gráfico sempre que a forma muda é um processo muito lento. Portanto, qualquer modelo que tenha tensores com formas dinâmicas, que mudam no tempo de execução, não é adequado para as TPUs.

Preenchimento

Programas de alto desempenho do Cloud TPU são aqueles com computação densa que pode ser facilmente dividida em blocos de 128 x 128. Quando uma computação de matriz não ocupa uma MXU inteira, o compilador preenche os tensores com zeros. O preenchimento apresenta duas desvantagens:

  • Os tensores preenchidos com zeros utilizam pouco o núcleo da TPU.
  • O preenchimento aumenta a quantidade de armazenamento em memória on-chip necessária para um tensor, o que pode resultar em erro de falta de memória em casos extremos.

O preenchimento é realizado automaticamente pelo compilador XLA quando necessário. No entanto, é possível determinar a quantidade de preenchimento por meio da ferramenta op_profile. Para evitar o preenchimento, escolha dimensões de tensor que sejam adequadas para as TPUs.

Dimensões

Escolher as dimensões de tensores adequadas ajuda muito para conseguir o máximo de desempenho do hardware da TPU, especialmente da MXU. O compilador XLA tenta usar o tamanho do lote ou a dimensão da característica para maximizar o uso da MXU. Portanto, um desses deve ser um múltiplo de 128. Caso contrário, o compilador preencherá um desses valores até chegar em 128. De preferência, o tamanho do lote e as dimensões da característica devem ser múltiplos de oito, permitindo, assim, conseguir um alto desempenho do subsistema de memória.

Operações

Veja a lista de operações do TensorFlow disponíveis.

Integração do 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 Visão geral do VPC Service Controls. 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.

Edge TPU

Os modelos de machine learning treinados na nuvem precisam executar cada vez mais inferências no nível mais avançado, ou seja, em dispositivos que operam no limite da Internet das coisas (IoT). Esses dispositivos incluem sensores e outros dispositivos inteligentes que reúnem dados em tempo real, tomam decisões inteligentes e, em seguida, realizam ações ou comunicam informações a outros dispositivos ou à nuvem.

Como esses dispositivos precisam funcionar com energia limitada, o Google desenvolveu o coprocessador Edge TPU para acelerar a inferência de ML em dispositivos de baixa potência. Um Edge TPU individual pode executar quatro trilhões de operações por segundo (4 TOPS) usando apenas 2 watts de energia. Em outras palavras, você recebe 2 TOPS por watt. Por exemplo, o Edge TPU pode executar modelos de última geração para dispositivos móveis, como o MobileNet V2, em quase 400 frames por segundo e com eficiência de consumo de energia.

Esse acelerador de ML de baixo consumo aumenta o Cloud TPU e o Cloud IoT para fornecer uma infraestrutura completa (nuvem para hardware, hardware e software) que facilita suas soluções baseadas em IA.

O Edge TPU está disponível para dispositivos de prototipagem e produção em vários formatos, incluindo um computador de placa única, um sistema em módulo, uma placa PCIe/M.2 e um módulo para superfície. Para mais informações sobre o Edge TPU e todos os produtos disponíveis, visite coral.ai.

A seguir