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. As TPUs foram desenvolvidas desde o início, aproveitando a extensa experiência e liderança do Google em aprendizado de máquina.

Com o Cloud TPU, é possível executar cargas de trabalho de aprendizado de máquina no hardware acelerador de TPUs de segunda geração 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

Os Cloud TPUs estão disponíveis nas seguintes zonas:

  • Estados Unidos (EUA)
    • us-central1-b
    • us-central1-c
    • us-central1-f (apenas o programa TFRC)
  • Europa (EU)
    • europe-west4-a
  • Ásia-Pacífico (APAC)
    • asia-east1-c

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. Isso significa que a compilação parcial de um modelo, em que a execução vai e volta entre o host e o dispositivo, usa o dispositivo de modo muito ineficiente porque ele fica ocioso na maior parte do tempo, à espera da chegada dos dados por meio do 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.

Estas são algumas das principais características do modelo de programação implementado pela TPUEstimator:

  • 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 Cloud TPU é amortizado pela execução de muitos passos de treinamento em loop.
  • Os dados de treinamento de entrada são transmitidos para uma fila de "alimentação" no Cloud TPU. Um programa em execução no Cloud TPU recupera lotes dessas filas durante cada passo do 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 passo da rede neural em todos os núcleos.

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. Normalmente, isso consiste nas operações de E/S que leem dados de treinamento, os pré-processamentos do código (por exemplo, decodificação de imagens compactadas, amostragem/recorte aleatório e montagem de minilotes de treinamento) e todas as partes de manutenção do gráfico, como o salvamento ou a restauração de pontos de verificação.

Práticas recomendadas para o desenvolvimento de modelos

O hardware do Cloud TPU contém uma unidade de matriz (MXU, na sigla em inglês) projetada para acelerar os programas que são dominados por multiplicações e convoluções de matrizes densas. O programas que não dedicam uma fração considerável do tempo de execução à realização de multiplicações de matrizes provavelmente não são adequados para o Cloud TPU. Por exemplo, um programa com computação que consiste preponderantemente de operações não matriciais, como "add", "reshape" ou "concatenate", provavelmente não atingirá uma alta utilização de MXU. A seguir há algumas diretrizes para ajudar você a escolher e criar modelos adequados para o Cloud TPU.

Réplicas

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. A TPUEstimator fornece um operador de gráfico para criar e executar uma computação replicada. Cada réplica é essencialmente uma cópia do gráfico de treinamento que é executado em cada núcleo. Essas cópias treinam minilotes com 1/8 do tamanho total do lote.

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 de memória on-chip necessária para um tensor, podendo 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 valores precisa 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.

Quando usar as 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, você pode decidir qual hardware é melhor para sua carga de trabalho com base nas seguintes diretrizes:

  • 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 que não estão ou não podem ser escritos em TensorFlow
    • 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 otimizados 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 que contêm 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.

É necessário que as cargas de trabalho de rede neural sejam capazes de executar várias iterações de todo o loop de treinamento na TPU. Ainda que esse não seja um requisito fundamental das TPUs, é uma das restrições atuais do ecossistema de software da TPU que são obrigatórias para aumentar a eficiência.

Próximas etapas

Esta página foi útil? Conte sua opinião sobre:

Enviar comentários sobre…