TPU v3

Neste documento, descrevemos a arquitetura e as configurações compatíveis do Cloud TPU v3.

arquitetura do sistema.

Cada chip de TPU v3 contém dois TensorCores. Cada TensorCore tem duas unidades de multiplicação de matriz (MXUs), uma unidade vetorial e uma unidade escalar. A tabela a seguir mostra as especificações das chaves e os valores correspondentes para um Pod de TPU v3.

Especificações da chave Valores de pod v3
Pico de computação por chip 123 teraflops (bf16)
Capacidade e largura de banda do HBM2 32 GiB e 900 GBps
Potência mínima/média/máxima medida 123/220/262 W
Tamanho do pod de TPU 1024 ícones
Topologia de interconexão Toro 2D
Pico de computação por pod 126 petaflops (bf16)
Largura de banda de redução total por pod 340 TB/s
Largura de banda de bissecção por pod 6,4 TB/s

O diagrama a seguir ilustra um chip TPU v3.

Diagrama de um chip TPU v3

Os detalhes arquitetônicos e as características de desempenho da TPU v3 estão disponíveis em Um supercomputador específico ao domínio para treinar redes neurais profundas (em inglês).

Benefícios de desempenho da TPU v3 em relação à v2

O aumento de FLOPS por TensorCore e a capacidade de memória nas configurações da TPU v3 podem melhorar o desempenho dos modelos das seguintes maneiras:

  • As configurações da TPU v3 oferecem benefícios de desempenho significativos de acordo com o TensorCore para modelos limitados por computação. Os modelos vinculados à memória nas configurações da TPU v2 podem não alcançar essa mesma melhoria de desempenho se também estiverem vinculados à memória nas configurações da TPU v3.

  • Nos casos em que os dados não cabem na memória nas configurações da TPU v2, a TPU v3 pode melhorar o desempenho e reduzir o novo cálculo de valores intermediários (rematerialização).

  • As configurações da TPU v3 executam novos modelos com tamanhos de lotes que não cabem nas configurações da TPU v2. Por exemplo, a TPU v3 pode permitir modelos ResNet mais profundos e imagens maiores com o RetinaNet.

Os modelos que são quase limitados à entrada ("alimentação") na TPU v2, porque as etapas de treinamento estão aguardando pela entrada, também podem ficar limitados à entrada no Cloud TPU v3. O guia de desempenho do pipeline pode ajudar você a resolver problemas de alimentação.

Configurações

Um Pod TPU v3 é composto por 1024 chips interconectados por links de alta velocidade. Para criar um dispositivo TPU v3 ou uma fração de pod, use a sinalização --accelerator-type no comando de criação da TPU (gcloud compute tpus tpu-vm). Especifique o tipo de acelerador especificando a versão e o número de núcleos de TPU. Por exemplo, para uma única TPU v3, use --accelerator-type=v3-8. Para uma fração do Pod v3 com 128 TensorCores, use --accelerator-type=v3-128.

O comando a seguir mostra como criar uma fração do Pod de TPU v3 com 128 TensorCores:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v3-128 \
    --version=tpu-vm-tf-2.16.1-pjrt

A tabela a seguir lista os tipos de TPU v3 compatíveis:

Versão da TPU Fim da compatibilidade
v3-8 (data final ainda não definida)
v3-32 (data final ainda não definida)
v3-128 (data final ainda não definida)
v3-256 (data final ainda não definida)
v3-512 (data final ainda não definida)
v3-1024 (data final ainda não definida)
v3-2048 (data final ainda não definida)

Para mais informações sobre como gerenciar TPUs, consulte Gerenciar TPUs. Para mais informações sobre a arquitetura do sistema do Cloud TPU, consulte Arquitetura do sistema.