As unidades de processamento de tensor (TPUs) são circuitos integrados de aplicativo específico (ASICs) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de machine learning (ML) e inteligência artificial (IA). Se você estiver treinando modelos de fundação complexos por semanas ou executando inferências em grande escala, as TPUs oferecem recursos de computação escalonáveis e especializados otimizados para estruturas de IA avançadas.
Este documento descreve a função das TPUs no Google Cloud, as especificações técnicas de cada versão de TPU disponível, as características de desempenho, as considerações de preços e como elas são mapeadas para as séries de máquinas do Compute Engine.
O que é uma TPU?
Uma TPU é um ASIC projetado sob medida para lidar com as demandas de operações de matrizes massivas essenciais aos algoritmos de ML.
Os principais componentes arquitetônicos de um sistema de TPU incluem:
TensorCore: a unidade de processamento de um chip de TPU. Cada TensorCore consiste em uma ou mais unidades de multiplicação de matrizes (MXUs) que usam arquitetura de matriz sistólica para realizar multiplicação de matrizes com alta eficiência, além de unidades escalares e vetoriais para fluxo de controle e computação geral.
SparseCore: processadores de fluxo de dados especializados projetados para acelerar operações esparsas. Os SparseCores trabalham com os TensorCores para processar grandes tabelas de embeddings de maneira eficiente, o que os torna ideais para acelerar modelos de recomendação e embeddings em grande escala.
Memória de alta largura de banda (HBM): grande memória física conectada ao chip da TPU que oferece largura de banda de memória enorme. A HBM permite o treinamento e a veiculação de modelos maiores com tamanhos de lote maiores.
Instância de TPU: uma instância de computação do Linux executada em um host de TPU e com acesso direto ao hardware de TPU, oferecendo acesso a bibliotecas de alto desempenho, conectividade SSH e registros de depuração de tempo de execução.
Frações e pods de TPU: um pod de TPU é um grande cluster contíguo de chips de TPU vinculados fisicamente. Uma fração de TPU é uma partição lógica de um pod (que inclui um ou mais hosts conectados usando interconexões de alta velocidade) alocada para executar uma única carga de trabalho.
Compilador XLA: o compilador de álgebra linear acelerada (XLA) compila gráficos de machine learning em instruções de máquina otimizadas que são executadas nos TensorCores da TPU.
Versões de TPU e séries de máquinas
O Compute Engine oferece suporte a várias gerações e versões de TPUs. A tabela a seguir fornece as principais especificações de cada versão de TPU compatível:
| Versão da TPU | Série de máquina | Pico de computação por chip (TFLOPs) | Memória e largura de banda da TPU | Núcleos por chip | Largura de banda de interconexão (ICI) por chip | Largura de banda da rede (DCN) por chip | Chips por pod |
|---|---|---|---|---|---|---|---|
| TPU7x (Ironwood) | tpu7x-standard |
BF16: 2.307 FP8: 4.614 |
192 GiB (7.380 GBps) |
2 TensorCores 4 SparseCores |
1.200 GBps | 100 Gbps | 9.216 |
| TPU v6e (Trillium) | ct6e-standard |
BF16: 918 FP8: 918 |
32 GiB (1.638 GBps) |
1 TensorCore 2 SparseCores |
800 GBps | 100 Gbps | 256 |
| TPU v5p | ct5p-hightpu |
BF16: 459 FP8: 459 |
95 GiB (2.765 GBps) |
2 TensorCores 4 SparseCores |
1.200 GBps | 50 Gbps | 8.960 |
Características de desempenho da TPU
As TPUs são otimizadas para executar operações densas de álgebra matricial com eficiência máxima. Para escolher quando uma TPU é a melhor opção para suas cargas de trabalho de ML, considere as diretrizes a seguir.
Casos de uso ideais e requisitos de carga de trabalho
A escolha do modelo de TPU certo depende das características computacionais, dos requisitos de memória e das necessidades de escalonabilidade da sua carga de trabalho. Selecione uma das guias a seguir para conferir as recomendações.
Treinamento com IA/ML
Pré-treinamento de modelos grandes: treinar modelos massivos do zero. Essas cargas de trabalho são vinculadas a computação e comunicação.
Ajuste e destilação de modelos: adaptação de modelos atuais usando métodos eficientes em termos de parâmetros ou treinamento de variantes menores de modelos.
Inferência de IA/ML
Veiculação de modelos on-line: implantação de modelos para interação em tempo real em que a baixa latência é essencial. Essas cargas de trabalho são limitadas pela largura de banda da memória durante a geração sequencial de tokens.
- Capacidade de TPU necessária: alta largura de banda de HBM e grande SRAM no chip para minimizar a latência de recuperação de dados.
- Versões recomendadas de TPU:
Execução da inferência em lote: processamento off-line de grandes conjuntos de dados em que o principal objetivo é a alta capacidade de processamento. Essas cargas de trabalho são vinculadas à computação durante a fase de pré-preenchimento do comando.
- Capacidade de TPU necessária: alta densidade de computação para maximizar a capacidade de processamento por dólar.
- Versões recomendadas de TPU:
Sistemas recomendadores
- Processamento de embeddings grandes: treinamento e veiculação de modelos de recomendação que usam tabelas de embeddings grandes. Essas cargas de trabalho são limitadas por capacidade
e comunicação.
- Capacidade de TPU necessária: hardware SparseCore especializado para processar operações esparsas em paralelo, grande capacidade de HBM e suporte para descarregamento na memória do host.
- Versões recomendadas de TPU:
Aprendizado por reforço
Execução de loops de aprendizado por reforço: gerenciamento de cargas de trabalho híbridas que exigem loops de geração de amostras e atualização de pesos de política.
Considerações sobre carga de trabalho
As TPUs podem não ser a melhor opção para as seguintes cargas de trabalho:
- Programas de álgebra linear que exigem ramificação lógica frequente ou contêm muitas operações algébricas de elemento.
- Cargas de trabalho que dependem de operadores personalizados em JAX ou PyTorch executados em CPUs.
- Cargas de trabalho científicas que exigem aritmética de usar pontos flutuantes de precisão dupla (FP64).
Considerações de preço
Os seguintes fatores determinam os preços da TPU no Google Cloud:
Versão da TPU
Local
Modelo de consumo
Para detalhes completos sobre preços, consulte a página Preços de TPU.
Modelos de consumo e compra
É possível provisionar instâncias de TPU usando qualquer uma das seguintes opções de consumo, dependendo dos seus requisitos de disponibilidade e tolerância a custos:
- Sob demanda: preços padrão de pagamento por utilização (PAYG) para execução imediata. Oferece flexibilidade máxima no ciclo de vida, mas a disponibilidade está sujeita a restrições de recursos físicos.
- Spot: capacidade excedente com descontos significativos (até 70%) de Google Cloud . Como o Google Cloud pode interromper ou excluir VMs spot para recuperar a capacidade com um aviso de 30 segundos, essas instâncias de computação são ideais para cargas de trabalho tolerantes a falhas e com pontos de verificação.
- Início flexível: permite solicitar instâncias de computação de um pool dedicado por até sete dias com maior disponibilidade.
- Reservas adiantadas: reservas que garantem capacidade e oferecem um preço com desconto em comparação com os preços sob demanda. Reserve a capacidade de TPU para uma data e hora futuras especificadas, por até 90 dias (modo de calendário) ou por um ano ou mais. Se oGoogle Cloud aprovar sua solicitação, você poderá começar a usar a capacidade durante o período da reserva.
Descontos por compromisso de uso (CUDs)
Ao solicitar uma reserva adiantada de um ano ou mais, você pode receber um CUD na capacidade de TPU reservada. Aplicar um CUD às suas instâncias de computação oferece uma redução significativa de preço em comparação com as taxas on demand em troca de um período de uso contínuo.
Para mais informações, consulte CUDs para o Compute Engine.
A seguir
- Analise as máquinas de TPU na família de máquinas com otimização de acelerador.
- Saiba como criar uma instância de TPU.