Google detalha TPUv4 e sua louca rede de IA opticamente reconfigurável

Aug 17, 2023

No Hot Chips 2023, o Google exibiu sua louca rede de IA opticamente reconfigurável. A empresa está trocando circuitos ópticos para obter melhor desempenho, menor consumo de energia e mais flexibilidade para seu cluster de treinamento de IA. A parte mais surpreendente é que eles já produzem isso há anos.

Isso está sendo feito ao vivo, então desculpe os erros de digitação.

O grande objetivo disso é unir os chips TPU do Google.

Aqui está o Google TPUv4 de 7 nm. Esperamos que esta semana comecemos a ouvir mais sobre TPUv5. O Google geralmente pode fazer artigos e apresentações sobre hardware antigo de uma geração. A TPU v4i era a versão de inferência, mas esta é mais a conversa focada na TPUv4.

O Google afirma que provisiona energia em excesso em comparação com a energia típica para que possa cumprir um SLA de tempo de serviço de 5 ms. Portanto, o TDP nos chips é muito mais alto, mas isso permite que o bursting atenda aos bursts de SLA.

Aqui está o diagrama da arquitetura TPUv4. O Google constrói esses chips TPU não apenas para serem um acelerador único, mas para serem expandidos e executados como parte de uma infraestrutura em grande escala.

Aqui estão as estatísticas do Google TPUv4 versus TPUv3 em uma das tabelas mais claras que já vimos sobre isso.

O Google mais que dobrou o pico de FLOPS, mas reduziu a potência entre TPUv3 e TPUv4.

O Google possui um acelerador SparseCore integrado ao TPUv4.

Aqui está o desempenho do TPUv4 SparseCore do Google.

A placa em si possui quatro chips TPUv4 e é refrigerada a líquido. O Google disse que teve que retrabalhar os data centers e as operações para mudar para a refrigeração líquida, mas a economia de energia valeu a pena. A válvula à direita controla o fluxo através dos tubos de resfriamento líquido. O Google diz que é como um controlador de velocidade do ventilador, mas para líquidos.

O Google também afirma que está usando PCIe Gen3 x16 de volta ao host, já que este era um design de 2020.

O Google tem energia entrando pela parte superior do rack, como muitos data centers, mas possui várias interconexões. Dentro de um rack, o Google pode usar DACs elétricos, mas fora de um rack, o Google precisa usar cabos ópticos.

Cada sistema possui 64 racks com 4.096 chips interconectados. De certa forma, os clusters de IA da NVIDIA em 256 nós têm metade do número de GPUs.

Também no final dos racks vemos um rack CDU. Se você quiser saber mais sobre resfriamento líquido, consulte Como funcionam os servidores de resfriamento líquido com Gigabyte e CoolIT. Teremos mais conteúdo de refrigeração líquida em breve. O Google diz que as taxas de fluxo de líquido são maiores do que a água na mangueira de um caminhão de bombeiros com gancho e escada.

Cada rack é um cubo 4x4x4 (64 nós) com comutação de circuitos ópticos (OCS) entre as TPUs. Dentro do rack, as conexões são DACs. As faces do cubo são todas ópticas.

Aqui está uma olhada no OCS. Em vez de usar um interruptor elétrico, o uso do OCS proporciona uma conexão direta entre os chips. O Google possui matrizes MEMS 2D internas, lentes, câmeras e muito mais. Evitar toda a sobrecarga da rede permite o compartilhamento de dados com mais eficiência. Como um aparte rápido, isso, em alguns aspectos, parece semelhante às TVs DLP.

O Google disse que tem mais de 16.000 conexões e distância de fibra suficiente no super pod para circundar o estado de Rhode Island.

Como há muita comunicação ponto a ponto, são necessários muitos fios de fibra.

Além disso, cada pool pode ser conectado a pools maiores.

O OCS, por ser reconfigurável, pode gerar maior utilização dos nós.

O Google pode então alterar as topologias ajustando o roteamento óptico.

Aqui o Google está mostrando os benefícios de diferentes topologias.

Isso é importante porque o Google afirma que as mudanças nas necessidades do modelo podem gerar mudanças no sistema.

Aqui está o dimensionamento do Google em escala logarítmica com acelerações lineares em até 3.072 chips.

O Google também aumentou a memória on-chip para 128 MB para manter o acesso aos dados local.

Aqui está a comparação do Google com o NVIDIA A100 com base no desempenho por watt.

Aqui está o modelo PaLM treinando em 6.144 TPUs em dois pods.

Esse é um número enorme!