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

Declaración del Problema

El ajuste fino y la implementación de modelos de lenguaje grandes (LLM) están cada vez más limitados por los costos de computación y la velocidad de iteración, más que por la disponibilidad de datos. Aunque las TPUs de Google Cloud ofrecen un alto rendimiento teórico, el ecosistema suele percibirse como que requiere un esfuerzo de ingeniería significativo en comparación con el maduro ecosistema de GPU basado en PyTorch. Este informe aborda la brecha en las herramientas abiertas documentando el proceso de extremo a extremo para portar una receta de entrenamiento nativa de GPU estándar (PyTorch, HuggingFace TRL, FSDP) a la pila de TPU basada en JAX (Tunix, Qwix, Orbax) para Gemma 4 31B. El estudio se centra específicamente en el Ajuste Fino Supervisado (SFT) con LoRA para la generación de código Verilog y la posterior implementación de inferencia, proporcionando una comparación empírica entre el hardware TPU y GPU bajo hiperparámetros idénticos.

Metodología

Los autores ejecutaron una comparación lado a lado del ajuste fino y la implementación de Gemma 4 31B en dos configuraciones de hardware distintas:

• Entrenamiento: TPU v5p-8 (4 chips) vs. 2× GPUs NVIDIA H100 de 80 GB.
• Inferencia: TPU v6e-8 (Trillium, 8 chips) vs. 2× GPUs NVIDIA H100 de 80 GB.
• Conjunto de datos: CodeV-R1 (10k pares de lenguaje natural a Verilog) para entrenamiento; NVlabs/verilog-eval (156 problemas de especificación a RTL) para evaluación.

Adaptaciones de Ingeniería (de GPU a TPU)

El núcleo de la metodología implicó un esfuerzo integral de portación para adaptar las tuberías de entrenamiento e inferencia de PyTorch a JAX:

1. Configuración de Framework y Malla: Transición de PyTorch FSDP a JAX con configuración explícita de malla de dispositivos (fsdp=1, tp=4 para v5p-8). Los autores notaron que tp=8 falló debido a que la cantidad de cabezas de atención de Gemma 4 no era divisible por 8.
2. Inyección de LoRA y Nomenclatura: Se abordaron las discrepancias en la nomenclatura de módulos entre HuggingFace (PyTorch) y Tunix (JAX). Notablemente, Tunix fusiona las proyecciones K y V en un único tensor kv_einsum, lo que requiere un manejo personalizado durante la fusión de puntos de control.
3. Fragmentación (Sharding) y Optimización: Se implementaron correcciones para las anotaciones de fragmentación de LoRA donde los parámetros inyectados por Qwix heredaron rangos de PartitionSpec desajustados de los pesos base. El optimizador se cambió a Optax AdamW, requiriendo lógica personalizada de fragmentación de estado.
4. Tubería de Datos: Se reemplazó el DataLoader de HuggingFace con la biblioteca Grain de Google, implementando lógica personalizada para el enmascaramiento de pérdidas (eliminación de indicaciones del sistema y cadenas de razonamiento) y el manejo de secuencias.
5. Puntos de Control y Conversión: Se desarrolló una tubería personalizada para fusionar adaptadores LoRA desde Orbax (formato JAX) de nuevo en Safetensors (formato PyTorch/HF) para inferencia, teniendo en cuenta las transposiciones de forma de tensor y las divisiones de kv_einsum fusionadas.
6. Configuración de Inferencia: Se desplegó vLLM-TPU en v6e-8 utilizando una imagen Docker específica (vllm/vllm-tpu:gemma4), navegando restricciones como la compilación de XLA al inicio, la asignación de HBM y la desactivación de entradas multimodales fragmentadas debido a las dimensiones de atención heterogéneas de Gemma 4.

Resultados Clave

Rendimiento de Entrenamiento

• Velocidad: La TPU v5p-8 completó el entrenamiento 1.61 veces más rápido (3.34 horas) que la línea base de 2×H100 (5.39 horas).
• Rendimiento: La TPU logró 763 tokens/seg frente a 486 tokens/seg en GPU.
• Costo: Debido a tarifas horarias más bajas (16.80 $ vs. 22.12 $) y una ejecución más rápida, el entrenamiento en TPU fue 2.12 veces más económico (56.11 $ vs. 119.23 $).
• Convergencia de Pérdida: Ambos modelos siguieron trayectorias de convergencia similares. La diferencia absoluta de pérdida (0.072 para TPU vs. 2.258 para GPU) se atribuyó a diferentes ámbitos de cálculo de pérdida (TPU usó pérdida solo de asistente; GPU usó pérdida de secuencia completa).

Rendimiento de Inferencia

• Contexto Corto (≤2048 tokens): La GPU mostró una ventaja marginal en el rendimiento máximo de salida (+6%) y una latencia de cola P99 ligeramente mejor, atribuido a un mayor ancho de banda de HBM por chip y la ausencia de sobrecarga de relleno de buckets XLA.
• Contexto Largo (≥4096 tokens): La TPU v6e-8 demostró dominio:
  • Rendimiento: +66% más de tokens de salida por segundo a 4096 tokens.
  • Tiempo hasta el Primer Token (TTFT): 23.6 veces más rápido a 4096 tokens (61 ms vs. 1,443 ms) y 7.1 veces más rápido a 8192 tokens.
  • Escalabilidad: La TPU sostuvo hasta 32 solicitudes concurrentes (QPS) en contexto medio, en comparación con la saturación de GPU en ~8, habilitado por la caché KV en fp8 (duplicando la capacidad efectiva) y una computación agregada superior.
• Costo: Para cargas de trabajo de contexto largo, la TPU fue 41% más económica por token de salida.

Calidad de Evaluación

En el punto de referencia de generación de Verilog, el modelo entrenado en GPU logró un pass@1 ligeramente superior (0.697 vs. 0.641) y un pass@5 (0.814 vs. 0.795). Los autores atribuyen esta diferencia marginal a variaciones en el objetivo de entrenamiento (enmascaramiento de pérdida) y varianza de muestreo estocástico, más que a capacidades de hardware.

Significancia y Afirmaciones

El artículo afirma proporcionar la primera demostración de extremo a extremo del ajuste fino y la implementación de Gemma 4 31B en hardware TPU, eliminando efectivamente una brecha crítica en el ecosistema de herramientas abiertas.

• Listo para Producción: El trabajo entrega una receta reproducible y lista para producción para desplegar Gemma 4 en infraestructura TPU, incluyendo configuraciones específicas de Docker, ajustes de malla y scripts de conversión de puntos de control.
• Costo Total de Propiedad (TCO): El estudio concluye que, para una carga de trabajo representativa de entrenamiento más implementación, la configuración TPU es 1.82 veces más económica que la línea base de GPU. Esta ventaja se impulsa principalmente por ahorros en costos de entrenamiento, con costos de inferencia que se acumulan significativamente para aplicaciones de contexto largo.
• Adecuación del Hardware: Los autores posicionan a la TPU como la infraestructura preferida para cargas de trabajo de contexto largo y pipelines de producción sensibles a costos, reconociendo al mismo tiempo las ventajas de la GPU en el rendimiento de contexto corto y la madurez del ecosistema (prototipado más sencillo).

El informe enfatiza que, aunque la sobrecarga de ingeniería para portar a JAX/Tunix no es trivial (estimada en aproximadamente una semana para configuración de malla, correcciones de fragmentación y reestructuración de tuberías), esto es un costo único amortizado a través de muchas ejecuciones de entrenamiento, haciendo de la TPU una alternativa viable y superior para la adaptación de LLM a gran escala.