TPU v6e

Este documento descreve a arquitetura e as configurações compatíveis do Cloud TPU v6e (Trillium).

O Trillium é o acelerador de IA de última geração da Cloud TPU. Em todas as plataformas técnicas, como a API e os registros, e ao longo deste documento, a Trillium será chamada de v6e.

Com uma pegada de 256 chips por pod, a v6e compartilha muitas semelhanças com a v5e. Esse sistema é otimizado para ser o produto de maior valor para treinamento, ajuste fino e veiculação de transformador, texto para imagem e rede neural convolucional (CNN, na sigla em inglês).

arquitetura do sistema.

Cada chip v6e contém um TensorCore. Cada TensorCore tem quatro unidades de multiplicação de matriz (MXU), uma unidade vetorial e uma escalar. A tabela a seguir mostra as principais especificações e os valores da TPU v6e em comparação com a TPU v5e.

Especificação v5e v6e
Desempenho/custo total de propriedade (TCO) (esperado) 0,65x 1
Pico de computação por chip (bf16) 197 TFLOPS 918 TFLOPs
Pico de computação por chip (Int8) 393 TOPs 1836 TOPs
Capacidade de HBM por chip 16 GB 32 GB
Largura de banda da HBM por chip 819 GBps 1.640 GBps
Largura de banda da interconexão entre chips (ICI) 1.600 Gbps 3.584 Gbps
Portas ICI por chip 4 4
DRAM por host 512 GiB 1536 GiB
Ícones por host 8 8
Tamanho do Pod de TPU 256 chips 256 chips
Topologia de interconexão Toroide 2D Toroide 2D
Computação máxima BF16 por pod 50,63 PFLOPs 234,9 PFLOPs
Largura de banda de redução total por pod 51,2 TB/s 102,4 TB/s
Largura de banda de bisection por pod 1,6 TB/s 3,2 TB/s
Configuração de NIC por host 2 NICs de 100 Gbps 4 x 200 Gbps NIC
Largura de banda da rede do data center por pod 6,4 Tbps 25,6 Tbps
Recursos especiais - SparseCore

Configurações aceitas

A TPU v6e oferece suporte ao treinamento de até 256 chips v6e e inferência de host único para até 8 chips.

A tabela a seguir mostra as formas de fatia 2D compatíveis com a v6e:

Topologia Chips do TPU Hosts VMs Tipo de acelerador (API TPU) Tipo de máquina (API GKE) Escopo Suporta inferência?
1x1 1 1/8 1 v6e-1 ct6e-standard-1t Subhost Sim
2x2 4 1/2 1 v6e-4 ct6e-standard-4t Subhost Sim
2x4 8 1 1 v6e-8 ct6e-standard-8t Host único Sim
2x4 8 1 2 - ct6e-standard-4t Host único Não
4x4 16 2 4 v6e-16 ct6e-standard-4t Vários hosts Não
4x8 32 4 8 v6e-32 ct6e-standard-4t Vários hosts Não
8x8 64 8 16 v6e-64 ct6e-standard-4t Vários hosts Não
8x16 128 16 32 v6e-128 ct6e-standard-4t Vários hosts Não
16x16 256 32 64 v6e-256 ct6e-standard-4t Vários hosts Não

As fatias com 8 chips (v6e-8) anexadas a uma única VM são otimizadas para inferência, permitindo que todos os 8 chips sejam usados em uma única carga de trabalho de fornecimento.

Para informações sobre o número de VMs para cada topologia, consulte Tipos de VM.

Tipos de VM

Cada VM de TPU v6e pode conter 1, 4 ou 8 chips. As fatias de 4 chips e menores têm o mesmo nó de acesso à memória não uniforme (NUMA). Para mais informações sobre nós NUMA, consulte Acesso à memória não uniforme na Wikipédia.

Diagrama de um host v6e

As fatias v6e são criadas usando VMs de meio host, cada uma com quatro chips de TPU. Há duas exceções a essa regra:

  • v6e-1: uma VM com apenas um chip, destinada principalmente a testes
  • v6e-8: uma VM de host completo otimizada para um caso de uso de inferência com todos os 8 chips conectados a uma única VM.

A tabela a seguir mostra uma comparação dos tipos de VM da TPU v6e:

Tipo de VM Número de vCPUs por VM RAM (GB) por VM Número de nós NUMA por VM
VM de 1 chip 44 176 1
VM de 4 chips 180 720 1
VM de 8 chips 180 1440 2

Especificar a configuração v6e

Ao alocar uma fatia de TPU v6e usando a API TPU, você especifica o tamanho e a forma usando os parâmetros AcceleratorType ou AcceleratorConfig.

Se você estiver usando o GKE, use a flag --machine-type para especificar um tipo de máquina compatível com a TPU que você quer usar. Para mais informações, consulte Implantar cargas de trabalho de TPU no GKE Standard na documentação do GKE.

Usar AcceleratorType

Ao alocar recursos da TPU, use AcceleratorType para especificar o número de TensorCores em uma fração. O valor especificado para AcceleratorType é uma string com o formato: v$VERSION-$TENSORCORE_COUNT. Por exemplo, v6e-8 especifica um slice de TPU v6e com 8 TensorCores.

O exemplo a seguir mostra como criar uma fatia de TPU v6e com 32 TensorCores usando AcceleratorType:

gcloud

    $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=zone \
        --accelerator-type=v6e-32 \
        --version=v2-alpha-tpuv6e

Console

  1. No console do Google Cloud, acesse a página TPUs:

    Acessar as TPUs

  2. Clique em Criar TPU.

  3. No campo Nome, insira um nome para o TPU.

  4. Na caixa Zona, selecione a zona em que você quer criar a TPU.

  5. Na caixa Tipo de TPU, selecione v6e-32.

  6. Na caixa Versão do software da TPU, selecione v2-alpha-tpuv6e. Ao criar uma VM do Cloud TPU, a versão do software da TPU especifica a versão do ambiente de execução da TPU que será instalada. Para mais informações, consulte Imagens de VM de TPU.

  7. Clique no botão Ativar fila.

  8. No campo Nome do recurso em fila, digite um nome para a solicitação de recurso em fila.

  9. Clique em Criar.

Usar AcceleratorConfig

Também é possível usar AcceleratorConfig para especificar o número de TensorCores que você quer usar. No entanto, como não há variantes de topologia 2D personalizadas para o TPU v6e, não há diferença entre o uso de AcceleratorConfig e AcceleratorType.

Para configurar uma TPU v6e usando AcceleratorConfig, use as flags --version e --topology. Defina --version como a versão da TPU que você quer usar e --topology como a disposição física dos chips da TPU na fração. O valor especificado para AcceleratorConfig é uma string com o formato AxB, em que A e B são as contagens de chips em cada direção.

O exemplo a seguir mostra como criar uma fração de TPU v6e com 32 TensorCores usando AcceleratorType organizados em uma topologia 4x8:

    $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=zone \
        --type=v6e \
        --topology=4x8 \
        --version=v2-alpha-tpuv6e