Google lança o Guia de Treinamento para Desenvolvedores do TPU Ironwood de sétima geração, detalhando a otimização de desempenho a nível de sistema

Notícias ME, 2 de abril (UTC+8), a Google anunciou recentemente um guia de treinamento para desenvolvedores voltado ao TPU Ironwood de sétima geração.
Este guia tem como objetivo ajudar os desenvolvedores a aproveitarem ao máximo o desempenho em nível de sistema do TPU Ironwood, para treinar e implantar modelos de IA de ponta de forma eficiente.
O TPU Ironwood é uma infraestrutura de IA personalizada, projetada para atender às demandas de poder computacional de modelos com trilhões de parâmetros, construído com tecnologias como interconexão entre chips (ICI), comutadores de luz (OCS), rede de data center (DCN) e memória de alta largura de banda agregada (HBM), formando um sistema completo que suporta até 9.216 chips.
O artigo detalha várias estratégias de otimização essenciais para esse hardware, incluindo:

• Utilizar a unidade de multiplicação de matrizes (MXU) nativa para suporte ao treinamento FP8, aumentando o throughput;
• Adotar a biblioteca de núcleos JAX otimizada para TPU, Tokamax, que processa tensores irregulares em modelos de contexto longo e modelos de especialistas mistos usando “atenção de respingo” e “multiplicação de matriz agrupada Megablox”;
• Utilizar o quarto geração de núcleos esparsos (SparseCore) para descarregar operações de comunicação coletiva e esconder latências;
• Ajustar finamente a alocação de SRAM rápida no chip do TPU (VMEM) para reduzir paradas de memória;
• E escolher a estratégia de particionamento ideal (como FSDP, TP, EP) com base na escala do modelo, arquitetura e comprimento da sequência.
  (Fonte: InFoQ)