TPU v6e

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

O Trillium é o mais recente acelerador de IA da Cloud TPU. Em todas as plataformas técnicas, como a API e os registros, e ao longo deste documento, o Trillium será chamado 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 e serviço de transformadores, conversão de texto em imagem e redes neurais convolucionais (CNNs).

arquitetura do sistema.

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

Configurações aceitas

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

As fatias com oito chips (v6e-8) anexadas a uma única VM são otimizadas para inferência, permitindo que todos os oito chips sejam usados em uma única carga de trabalho de serviço. É possível realizar inferências de vários hosts usando o Pathways on Cloud. Para mais informações, consulte Realizar inferência multihost usando o Pathways

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

Tipos de VM

Cada VM de TPU v6e pode conter 1, 4 ou 8 chips. As partições de quatro chips ou menos 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.

As frações 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 oito chips conectados a uma única VM.

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

Especificar a configuração do v6e

Ao alocar uma fração de TPU v6e usando a API TPU, especifique o tamanho e a forma usando o parâmetro AcceleratorType.

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 Planejar TPUs no GKE na documentação do GKE.

Usar AcceleratorType

Ao alocar recursos de 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 fração de TPU v6e com 32 TensorCores usando AcceleratorType:

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

A seguir

• Executar treinamento e inferência usando a TPU v6e