No início deste ano, o Google lançou sua sexta geração e até agora a TPU mais poderosa, o Trillium. Hoje, o Trillium está oficialmente disponível para clientes do Google Cloud.
O Google usou a TPU Trillium para treinar o Gemini 2.0, seu modelo de IA mais poderoso até hoje. Agora, empresas e startups podem aproveitar a mesma infraestrutura poderosa, eficiente e sustentável.
O núcleo do supercomputador de IA: TPU Trillium
A TPU Trillium é um componente essencial do Google Cloud AI Hypercomputer. O AI Hypercomputer é uma arquitetura de supercomputador inovadora que integra hardware otimizado para desempenho, software aberto, frameworks líderes de ML e um sistema de modelo de consumo flexível. Com o lançamento oficial da TPU Trillium, o Google também fez melhorias importantes na camada de software aberto do AI Hypercomputer, incluindo um compilador XLA otimizado e frameworks populares como JAX, PyTorch e TensorFlow, para atingir a melhor relação custo-benefício em treinamento, ajuste e serviço de IA.
Além disso, recursos como a descarga do host usando DRAM de host em grande escala (complementando a memória de alta largura de banda ou HBM) fornecem um nível mais alto de eficiência. O AI Hypercomputer permite que você extraia o máximo valor de implantações de mais de 100.000 chips Trillium de cada arquitetura de rede Jupiter, que possui largura de banda bidirecional de 13 Petabits/segundo, capaz de escalar um único trabalho de treinamento distribuído para centenas de milhares de aceleradores.
Clientes como a AI21Labs já estão usando o Trillium para entregar soluções de IA significativas para seus clientes mais rapidamente:
Barak Lenz, CTO da AI21Labs, afirma: “Na AI21, estamos constantemente buscando melhorar o desempenho e a eficiência dos nossos modelos de linguagem Mamba e Jamba. Como usuários de longa data da TPU v4, ficamos impressionados com os recursos do Trillium do Google Cloud. Os avanços em escala, velocidade e eficiência de custo são bastante significativos. Acreditamos que o Trillium desempenhará um papel crucial na aceleração do desenvolvimento de nossos modelos de linguagem complexos de próxima geração, permitindo-nos fornecer soluções de IA mais poderosas e acessíveis aos nossos clientes.”
Desempenho significativamente aprimorado do Trillium, com vários indicadores quebrando recordes
Em comparação com a geração anterior, o Trillium apresenta melhorias significativas em:
Desempenho de treinamento mais de 4 vezes maior
Aumento de 3 vezes na taxa de transferência de inferência
Melhoria de 67% na eficiência energética
Aumento de 4,7 vezes no desempenho de computação de pico por chip
Capacidade de memória de alta largura de banda (HBM) duplicada
Largura de banda de interconexão entre chips (ICI) duplicada
100.000 chips Trillium em uma única arquitetura de rede Jupiter
Desempenho de treinamento por dólar 2,5 vezes maior, desempenho de inferência por dólar 1,4 vezes maior
Esses aprimoramentos permitem que o Trillium se destaque em várias cargas de trabalho de IA, incluindo:
Escalonamento de cargas de trabalho de treinamento de IA
Treinamento de LLMs, incluindo modelos densos e modelos de especialistas mistos (MoE)
Desempenho de inferência e agendamento de ensembles
Modelos densos de incorporação
Fornecimento de relação custo-benefício para treinamento e inferência
Como o Trillium se destaca em diferentes cargas de trabalho?
Escalonamento de cargas de trabalho de treinamento de IA
Treinar modelos grandes como o Gemini 2.0 requer uma grande quantidade de dados e computação. A capacidade de escalonamento quase linear do Trillium permite que esses modelos acelerem significativamente o treinamento, distribuindo eficientemente as cargas de trabalho em várias máquinas Trillium conectadas por interconexões de chips de alta velocidade em pods de 256 chips e nossa rede de data center Jupiter de última geração. Isso é alcançado por meio de várias TPU, uma tecnologia de pilha completa para treinamento em larga escala, e ainda mais otimizado pelo Titanium, um sistema de descarga dinâmico em nível de data center que abrange desde adaptadores de host até arquitetura de rede.
O Trillium alcançou 99% de eficiência de escalonamento em uma implantação de 12 pods com 3072 chips e mostrou 94% de eficiência de escalonamento em 24 pods com 6144 chips para pré-treinar o gpt3-175b, mesmo quando executado em uma rede de data center.
Treinamento de LLMs, incluindo modelos densos e modelos de especialistas mistos (MoE)
LLMs como o Gemini são intrinsecamente poderosos e complexos, com bilhões de parâmetros. Treinar esses LLMs densos requer uma enorme capacidade de computação e otimizações de software projetadas em conjunto. O Trillium é 4 vezes mais rápido que a geração anterior de Cloud TPU v5e no treinamento de LLMs densos como Llama-2-70b e gpt3-175b.
Além dos LLMs densos, o treinamento de LLMs usando arquiteturas de especialistas mistos (MoE) é uma abordagem cada vez mais popular que combina várias redes neurais "especialistas", cada uma especializada em diferentes aspectos de uma tarefa de IA. Gerenciar e coordenar esses especialistas durante o treinamento adiciona complexidade em comparação com o treinamento de um único modelo monolítico. O Trillium é 3,8 vezes mais rápido que a geração anterior de Cloud TPU v5e no treinamento de modelos MoE.
Além disso, em comparação com a Cloud TPU v5e, a TPU Trillium oferece 3 vezes mais memória de acesso aleatório dinâmica (DRAM) do host. Isso descarrega alguns cálculos para o host, ajudando a maximizar o desempenho em larga escala e um bom throughput. O recurso de descarga do host do Trillium forneceu mais de 50% de aumento de desempenho em termos de utilização de FLOP do modelo (MFU) no treinamento do modelo Llama-3.1-405B.
Desempenho de inferência e agendamento de ensembles
Durante a inferência, a importância da inferência de várias etapas está crescendo, exigindo que os aceleradores processem eficientemente o aumento das necessidades computacionais. O Trillium oferece grandes avanços para cargas de trabalho de inferência, permitindo a implantação de modelos de IA mais rápida e eficiente. Na verdade, o Trillium oferece nosso melhor desempenho de inferência de TPU para modelos de difusão de imagem e LLMs densos. Nossos testes mostram que a taxa de transferência de inferência relativa (imagens por segundo) do Stable Diffusion XL é mais de 3 vezes maior em comparação com a Cloud TPU v5e, enquanto a taxa de transferência de inferência relativa (tokens por segundo) do Llama2-70B é quase 2 vezes maior.
O Trillium é nossa TPU de maior desempenho para casos de uso de inferência offline e de servidor. O gráfico abaixo mostra que a taxa de transferência de inferência relativa offline (imagens por segundo) do Stable Diffusion XL é 3,1 vezes maior e a taxa de transferência de inferência relativa do servidor é 2,9 vezes maior em comparação com a Cloud TPU v5e.
Além do melhor desempenho, o Trillium introduziu novos recursos de agendamento de ensembles. Este recurso permite que o sistema de agendamento do Google tome decisões inteligentes de agendamento de trabalhos para melhorar a disponibilidade e a eficiência geral das cargas de trabalho de inferência quando vários réplicas estão presentes em um ensemble. Ele fornece uma maneira de gerenciar vários fragmentos de TPU que executam cargas de trabalho de inferência de host único ou vários hosts, incluindo por meio do Google Kubernetes Engine (GKE). Agrupar esses fragmentos em um ensemble permite ajustar facilmente o número de réplicas para corresponder à demanda.
Modelos densos de incorporação
Com a adição do SparseCore de terceira geração, o Trillium melhorou o desempenho dos modelos densos de incorporação em 2 vezes e o desempenho do DLRM DCNv2 em 5 vezes.
O SparseCore é um processador de fluxo de dados que fornece uma base de arquitetura mais adaptável para cargas de trabalho densas de incorporação. O SparseCore de terceira geração do Trillium é excelente na aceleração de operações dinâmicas e dependentes de dados, como coleta dispersa, soma de segmentos esparsos e particionamento.
Fornecimento de relação custo-benefício para treinamento e inferência
Além do desempenho e da escala absolutos necessários para treinar algumas das maiores cargas de trabalho de IA do mundo, o Trillium também visa otimizar o desempenho por dólar. Até o momento, o Trillium oferece um desempenho por dólar 2,1 vezes maior que a Cloud TPU v5e e 2,5 vezes maior que a Cloud TPU v5p no treinamento de LLMs densos como Llama2-70b e Llama3.1-405b.
O Trillium é excelente no processamento paralelo de modelos grandes de forma econômica. Ele foi projetado para permitir que pesquisadores e desenvolvedores forneçam modelos de imagem poderosos e eficientes a um custo muito menor do que antes. O custo de gerar mil imagens no Trillium é 27% menor para inferência offline e 22% menor para inferência de servidor no SDXL em comparação com a Cloud TPU v5e.
Elevando a inovação em IA para um novo nível
O Trillium representa um grande avanço na infraestrutura de IA do Google Cloud, oferecendo desempenho, escalabilidade e eficiência incríveis para uma variedade de cargas de trabalho de IA. Com sua capacidade de escalar para centenas de milhares de chips usando software de co-design de classe mundial, o Trillium permite que você realize avanços mais rápidos e entregue soluções de IA excepcionais. Além disso, a excelente relação custo-benefício do Trillium o torna uma opção econômica para organizações que desejam maximizar o valor do investimento em IA. À medida que o cenário da IA continua a evoluir, o Trillium demonstra o compromisso do Google Cloud em fornecer infraestrutura de ponta para ajudar as empresas a liberar todo o potencial da IA.
Introdução oficial: https://cloud.google.com/blog/products/compute/trillium-tpu-is-ga