Resumo Técnico: TPU vs. GPU para LoRA SFT do Gemma 4 31B

Declaração do Problema

O ajuste fino e a inferência de modelos de linguagem de grande escala (LLMs) estão cada vez mais limitados pelos custos de computação e pelas velocidades de iteração, e não pela disponibilidade de dados. Embora os TPUs do Google Cloud ofereçam alto desempenho teórico, o ecossistema é frequentemente percebido como exigindo um esforço de engenharia significativo em comparação com o ecossistema maduro de GPUs baseado em PyTorch. Este relatório aborda a lacuna nas ferramentas abertas documentando o processo de ponta a ponta da portabilidade de uma receita padrão de treinamento nativa de GPU (PyTorch, HuggingFace TRL, FSDP) para a pilha baseada em JAX do TPU (Tunix, Qwix, Orbax) para o Gemma 4 31B. O estudo visa especificamente o Ajuste Fino Supervisionado (SFT) com LoRA para geração de código Verilog e subsequente serviço de inferência, fornecendo uma comparação empírica entre hardware TPU e GPU sob hiperparâmetros idênticos.

Metodologia

Os autores executaram uma comparação lado a lado do ajuste fino e da inferência do Gemma 4 31B em duas configurações de hardware distintas:

• Treinamento: TPU v5p-8 (4 chips) vs. 2× GPUs NVIDIA H100 80GB.
• Inferência: TPU v6e-8 (Trillium, 8 chips) vs. 2× GPUs NVIDIA H100 80GB.
• Conjunto de Dados: CodeV-R1 (10 mil pares de linguagem natural para Verilog) para treinamento; NVlabs/verilog-eval (156 problemas de especificação para RTL) para avaliação.

Adaptações de Engenharia (GPU para TPU)

O núcleo da metodologia envolveu um esforço abrangente de portabilidade para adaptar os pipelines de treinamento e inferência de PyTorch para JAX:

1. Configuração de Framework e Malha: Transição do PyTorch FSDP para JAX com configuração explícita de malha de dispositivos (fsdp=1, tp=4 para v5p-8). Os autores observaram que tp=8 falhou devido ao número de cabeças de atenção do Gemma 4 não ser divisível por 8.
2. Injeção e Nomeação LoRA: Resolução de discrepâncias entre a nomeação de módulos do HuggingFace (PyTorch) e do Tunix (JAX). Notavelmente, o Tunix funde as projeções K e V em um único tensor kv_einsum, exigindo tratamento personalizado durante a mesclagem de checkpoints.
3. Fragmentação (Sharding) e Otimização: Implementação de correções para anotações de fragmentação LoRA onde parâmetros injetados pelo Qwix herdaram ranks PartitionSpec incompatíveis dos pesos base. O otimizador foi alterado para Optax AdamW, exigindo lógica personalizada de fragmentação de estado.
4. Pipeline de Dados: Substituição do DataLoader do HuggingFace pela biblioteca Grain do Google, implementando lógica personalizada para mascaramento de perda (removendo prompts de sistema e cadeias de raciocínio) e manipulação de sequências.
5. Checkpointing e Conversão: Desenvolvimento de um pipeline personalizado para mesclar adaptadores LoRA do Orbax (formato JAX) de volta ao formato Safetensors (PyTorch/HF) para inferência, considerando transposições de forma de tensor e divisões de kv_einsum fundidas.
6. Configuração de Inferência: Implantação do vLLM-TPU no v6e-8 usando uma imagem Docker específica (vllm/vllm-tpu:gemma4), navegando por restrições como compilação XLA na inicialização, alocação de HBM e a desativação de entradas multimodais fragmentadas devido às dimensões de atenção heterogêneas do Gemma 4.

Principais Resultados

Desempenho de Treinamento

• Velocidade: O TPU v5p-8 concluiu o treinamento 1,61× mais rápido (3,34 horas) que a linha de base de 2×H100 (5,39 horas).
• Vazão: O TPU alcançou 763 tokens/seg vs. 486 tokens/seg na GPU.
• Custo: Devido a taxas horárias mais baixas ($16,80 vs. $22,12) e execução mais rápida, o treinamento em TPU foi 2,12× mais barato ($56,11 vs. $119,23).
• Convergência de Perda: Ambos os modelos seguiram trajetórias de convergência semelhantes. A diferença absoluta de perda (0,072 para TPU vs. 2,258 para GPU) foi atribuída a diferentes escopos de cálculo de perda (TPU usou perda apenas do assistente; GPU usou perda de sequência completa).

Desempenho de Inferência

• Contexto Curto (≤2048 tokens): A GPU mostrou uma vantagem marginal na vazão máxima de saída (+6%) e latência de cauda P99 ligeiramente melhor, atribuída à maior largura de banda de HBM por chip e à ausência de sobrecarga de preenchimento de buckets XLA.
• Contexto Longo (≥4096 tokens): O TPU v6e-8 demonstrou domínio:
  • Vazão: +66% mais tokens de saída/seg em 4096 tokens.
  • Tempo até o Primeiro Token (TTFT): 23,6× mais rápido em 4096 tokens (61 ms vs. 1.443 ms) e 7,1× mais rápido em 8192 tokens.
  • Escalabilidade: O TPU sustentou até 32 solicitações simultâneas (QPS) em contexto médio, comparado à saturação da GPU em ~8, habilitado por cache KV fp8 (dobrando a capacidade efetiva) e maior computação agregada.
• Custo: Para cargas de trabalho de contexto longo, o TPU foi 41% mais barato por token de saída.

Qualidade de Avaliação

No benchmark de geração de Verilog, o modelo treinado na GPU alcançou um pass@1 ligeiramente superior (0,697 vs. 0,641) e pass@5 (0,814 vs. 0,795). Os autores atribuem essa diferença marginal a variações no objetivo de treinamento (mascaramento de perda) e variância de amostragem estocástica, e não a capacidades de hardware.

Significado e Afirmações

O artigo afirma fornecer a primeira demonstração de ponta a ponta do ajuste fino e serviço do Gemma 4 31B em hardware TPU, removendo efetivamente uma lacuna crítica no ecossistema de ferramentas abertas.

• Prontidão para Produção: O trabalho entrega uma receita reproduzível e pronta para produção para implantar o Gemma 4 em infraestrutura TPU, incluindo configurações específicas de Docker, configurações de malha e scripts de conversão de checkpoint.
• Custo Total de Propriedade (TCO): O estudo conclui que, para uma carga de trabalho representativa de treinamento mais serviço, a configuração TPU é 1,82× mais barata que a linha de base de GPU. Essa vantagem é impulsionada principalmente pelas economias de custo de treinamento, com os custos de inferência se acumulando significativamente para aplicações de contexto longo.
• Adequação de Hardware: Os autores posicionam o TPU como a infraestrutura preferida para cargas de trabalho de contexto longo e pipelines de produção sensíveis a custos, enquanto reconhecem as vantagens da GPU em vazão de contexto curto e maturidade do ecossistema (prototipagem mais simples).

O relatório enfatiza que, embora o sobrecarga de engenharia da portabilidade para JAX/Tunix seja não trivial (estimada em cerca de uma semana para configuração de malha, correções de fragmentação e reestruturação de pipeline), este é um custo único amortizado em várias execuções de treinamento, tornando o TPU uma alternativa viável e superior para adaptação de LLMs em grande escala.