Google integra fuzzing en el autotuner de XLA y detecta 142 bugs silenciosos del compilador en TPU

Si un compilador falla con un error, lo notas. Si produce números incorrectos en silencio, te enteras tras una semana de entrenamiento, cuando la pérdida (loss) se comporta de forma extraña y la factura del clúster ya ha llegado. Precisamente estos bugs silenciosos analiza el equipo de Google en su trabajo XLA from the Trenches, y su solución es sorprendentemente pragmática.

Por qué los bugs silenciosos son peores que los fallos

XLA convierte el grafo de cómputo de un modelo en código máquina para TPU. En el camino fusiona operaciones en kernels grandes (fusion), distribuye los datos en registros vectoriales y memoria, y aplica decenas de optimizaciones más. Cada una puede ser correcta por separado y fallar en una combinación poco frecuente de formas, tipos y parámetros.

Las pruebas estáticas del CI comprueban casos ideados de antemano. Ni siquiera millones de pruebas así cubren todas las combinaciones que surgen en modelos reales. Además, el bug no se manifiesta como un fallo: todo se ejecuta, todo es rápido, solo que el resultado es ligeramente distinto. En el entrenamiento de un modelo grande, el coste de un error así se mide en millones de dólares.

Cómo funciona

La idea clave es no construir una infraestructura de verificación aparte y costosa, sino engancharse a algo que ya funciona. El autotuner de rendimiento ya prueba cada día distintas variantes de compilación del mismo kernel en busca de la más rápida. Los autores añadieron a este ciclo una comparación de resultados: la salida de cada variante se contrasta con la del kernel fusionado por defecto. Si hay discrepancia, alguna de las variantes se está compilando mal.

Si la sospecha recae en un bug dentro del grafo, el sistema compara además el resultado con una referencia sin fusionar, donde las operaciones se ejecutan por separado. Así se puede saber si el culpable es precisamente la fusión.

A continuación entra en juego un aislador en dos etapas. Dividir simplemente el kernel en partes suele enmascarar el bug, porque el problema vive justamente en la versión fusionada. Por eso el aislador lee paso a paso los registros vectoriales de hardware dentro del propio kernel fusionado y encuentra la instrucción en la que los valores divergen por primera vez. Además, los autores emplean generación de programas sintéticos y un enjambre de agentes de IA que escriben parches para los problemas encontrados.

Qué encontraron

En los clústeres de TPU de Google, el sistema ayudó a corregir más de 142 bugs que el CI había pasado por alto. La distribución resulta curiosa: el 43% no estaba en las transformaciones de alto nivel del grafo, sino en los generadores de código de bajo nivel para registros vectoriales y memoria. Es decir, se rompía justo donde las pruebas habituales a nivel de grafo apenas miran.

Limitaciones y conclusiones

Para ser justos, el trabajo tiene puntos débiles. Los 142 bugs se contabilizan con una sola cifra, sin ablaciones por componente, así que no está claro qué parte del sistema aporta más: el fuzzing, el aislador o los agentes. Las condiciones de carrera de memoria no deterministas, los auténticos heisenbugs, quedan fuera del alcance del método. Además, comparar con el kernel por defecto presupone que este es correcto, y no siempre lo es.

La conclusión práctica va más allá de TPU. Si tienes un autotuner, un compilador JIT o cualquier mecanismo que recorra implementaciones equivalentes de una misma operación, ya estás pagando por calcular varias respuestas sobre las mismas entradas. Añadir una comparación de resultados con tolerancia al ruido numérico es barato, y detecta precisamente los errores que nadie vería. Para los equipos que entrenan modelos grandes sobre su propio stack, es una de las formas más económicas de atrapar los defectos antes de que devoren el presupuesto.