TPU7x (Ironwood)

Esta página descreve a arquitetura e as configurações disponíveis para a TPU7x, a mais recente TPU disponível no Google Cloud. A TPU7x é a primeira versão da família Ironwood, Google Cloud's sétima geração de TPU. A geração Ironwood foi projetada para treinamento e inferência de IA em grande escala.

Com 9.216 chips por pod, a TPU7x compartilha muitas semelhanças com TPU v5p. A TPU7x oferece alto desempenho para modelos densos e MoE em grande escala, pré-treinamento, amostragem e inferência com uso intenso de decodificação.

Para usar a TPU7x, é necessário usar o Google Kubernetes Engine (GKE). Para mais informações, consulte Sobre TPUs no GKE.

Também é possível usar a TPU7x e o GKE com o TPU Cluster Director. O TPU Cluster Director está disponível por meio de uma reserva no modo "All Capacity", que oferece acesso total a toda a capacidade reservada (sem retenções) e visibilidade completa da topologia de hardware da TPU, do status de utilização e do status de integridade. Para mais informações, consulte Visão geral do modo "All Capacity".

Para ter acesso à TPU7x, entre em contato com a equipe de conta.

Arquitetura do sistema

Cada chip da TPU7x contém dois TensorCores e quatro SparseCores. A tabela a seguir mostra as principais especificações e os valores da TPU7x em comparação com as gerações anteriores.

Especificação v5p v6e (Trillium) TPU7x (Ironwood)
Número de chips por pod 8960 256 9216
Pico de computação por chip (BF16) (TFLOPs) 459 918 2307
Pico de computação por chip (FP8) (TFLOPs) 459 918 4614
Capacidade de HBM por chip (GiB) 95 32 192
Largura de banda de HBM por chip (GiBps) 2575 1638 7380
Número de vCPUs (VM de quatro chips) 208 180 224
RAM (GB) (VM de quatro chips) 448 720 960
Número de TensorCores por chip 2 1 2
Número de SparseCores por chip 4 2 4
Largura de banda bidirecional da interconexão entre chips (ICI) por chip (GBps) 1200 800 1200
Largura de banda da rede do data center (DCN) por chip (Gbps) 50 100 100

O diagrama a seguir ilustra a arquitetura do Ironwood:

Diagrama da arquitetura do Ironwood

Hierarquia de memória

A TPU7x apresenta um sistema de memória de vários níveis, e o gerenciamento da movimentação de dados entre esses níveis é fundamental para o desempenho:

  • Memória de alta largura de banda (HBM): cada chip está equipado com 192 GB de HBM, com largura de banda de aproximadamente 7,37 TB/s. A grande capacidade de HBM permite tamanhos de lote grandes, o que pode melhorar a capacidade de processamento. No entanto, apesar do tamanho, a HBM ainda pode ser um gargalo, principalmente para operações vetoriais vinculadas à memória ou padrões de acesso aos dados ineficientes.
  • Memória vetorial (VMEM): a VMEM é uma SRAM menor no chip (memória de acesso aleatório estática) com largura de banda significativamente maior para a unidade de multiplicação de matrizes (MXU) do que a HBM. Essa memória atua como um bloco de rascunho de alta velocidade para kernels personalizados. O tamanho desse buffer é um parâmetro ajustável. A otimização do tamanho do buffer é fundamental para ajustar os kernels Pallas personalizados, já que os tamanhos de bloco geralmente são restritos pela VMEM disponível.
  • Memória do host e PCIe:cada conjunto de quatro chips de TPU está conectado a um host de CPU usando uma rede PCIe. Embora essa conexão tenha uma largura de banda muito menor do que a HBM, a memória principal do host pode ser usada para descarregar ativações ou estados do otimizador para liberar a HBM, uma técnica particularmente útil para gerenciar a pressão de memória em modelos grandes.

Para mais informações sobre como gerenciar com eficiência a movimentação de dados entre os níveis da hierarquia de memória da TPU7x, consulte Otimizações de desempenho do Ironwood.

Arquitetura de chiplet duplo

O modelo de programação do Ironwood permite acessar dois chiplets de TPU em vez de uma única arquitetura de núcleo lógico (também conhecida como MegaCore) usada em gerações anteriores (TPU v4 e v5p). Essa mudança melhora a relação custo-benefício e a eficiência da fabricação do chip. Embora isso represente uma mudança arquitetônica, o novo design garante que você possa reutilizar modelos de software atuais com mudanças mínimas.

As TPUs do Ironwood são compostas por dois chiplets distintos, cada um com seu próprio espaço de memória dedicado. Essa é uma mudança em relação ao espaço de memória unificado da arquitetura MegaCore.

  • Composição do chiplet: cada chiplet é uma unidade independente com um TensorCore, dois SparseCores e 96 GB de memória de alta largura de banda (HBM).

  • Interconexão de alta velocidade: os dois chiplets são conectados por uma interface de chip a chip (D2D) que é seis vezes mais rápida do que um link de interconexão entre chips (ICI) 1D. A comunicação entre chiplets é gerenciada usando operações coletivas.

Modelo de programação e exposição de framework

O modelo de programação do Ironwood é semelhante ao das gerações de TPU anteriores à v4, como a TPU v3. A nova arquitetura é exposta das seguintes maneiras:

  • Dois dispositivos por chip:frameworks como o JAX expõem cada chip do Ironwood como dois "dispositivos" separados, um para cada chiplet.

  • Especificação do chiplet:é possível especificar qual chiplet usar para uma computação. O JAX adiciona uma quarta dimensão à especificação de topologia para distinguir entre chiplets. Esse design permite reutilizar modelos de software atuais com mudanças mínimas.

Para mais informações sobre como alcançar o desempenho ideal com a arquitetura de chiplet duplo, consulte Recomendações de desempenho para a arquitetura de chiplet duplo do Ironwood

Configurações aceitas

Os chips da TPU7x têm uma conexão direta com os chips vizinhos mais próximos em três dimensões, o que resulta em uma malha 3D de conexões de rede. As frações maiores que 64 chips são compostas por um ou mais "cubos" de chips de 4x4x4.

Os chips da TPU7x têm uma topologia de interconexão de resíduo 3D. Essa topologia permite que as frações sejam escalonar verticalmente verticalmente para até 9.216 chips. Ela tem largura de banda bidirecional de 200 GBps por eixo para comunicação entre chips em um pod.

A tabela a seguir mostra as formas de fração 3D comuns que são compatíveis com a TPU7x:

Topologia Chips de TPU Hosts VMs Cubos Escopo
2x2x1 4 1 1 1/16 Host único
2x2x2 8 2 2 1/8 Vários hosts
2x2x4 16 4 4 1/4 Vários hosts
2x4x4 32 8 8 1/2 Vários hosts
4x4x4 64 16 16 1 Vários hosts
4x4x8 128 32 32 2 Vários hosts
4x8x8 256 64 64 4 Vários hosts
8x8x8 512 128 128 8 Vários hosts
8x8x16 1024 256 256 16 Vários hosts
8x16x16 2048 512 512 32 Vários hosts

VM da TPU7x

Cada máquina virtual (VM) da TPU7x contém quatro chips. Cada VM tem acesso a dois nós NUMA. Para mais informações sobre nós NUMA, consulte Acesso à memória não uniforme na Wikipédia.

Todas as frações da TPU7x usam VMs de host completo de quatro chips. As especificações técnicas de uma VM da TPU7x são:

  • Número de vCPUs por VM: 224
  • RAM por VM: 960 GB
  • Número de nós NUMA por VM: 2

Hyperdisk

Por padrão, o disco de inicialização da VM para a TPU7x é o Hyperdisk Balanced. É possível anexar mais discos à VM de TPU para armazenamento adicional. Os seguintes tipos de disco são compatíveis com a TPU7x:

  • Hyperdisk Balanced
  • Hyperdisk ML

Para mais informações sobre o Hyperdisk, consulte Visão geral do Hyperdisk. Para mais informações sobre as opções de armazenamento do Cloud TPU, consulte Opções de armazenamento para dados do Cloud TPU.

A seguir