Google integra fuzzing ao autotuner do XLA e encontra 142 bugs silenciosos do compilador em TPU

Se um compilador falha com um erro, você percebe. Se ele produz números errados em silêncio, você descobre só após uma semana de treinamento, quando a loss se comporta de forma estranha e a conta do cluster já chegou. É exatamente esse tipo de bug silencioso que a equipe do Google analisa no trabalho XLA from the Trenches, e a solução deles é surpreendentemente pé no chão.

Por que bugs silenciosos são piores que falhas

O XLA converte o grafo de computação de um modelo em código de máquina para TPU. No caminho, ele funde operações em kernels grandes (fusion), distribui os dados em registradores vetoriais e memória e aplica dezenas de outras otimizações. Cada uma pode estar correta isoladamente e falhar em uma combinação rara de formas, tipos e parâmetros.

Os testes estáticos do CI verificam casos imaginados de antemão. Nem milhões desses testes cobrem todas as combinações que surgem em modelos reais. Além disso, o bug não se manifesta como falha: tudo executa, tudo é rápido, só que o resultado está um pouco diferente. No treinamento de um modelo grande, o custo de um erro assim é medido em milhões de dólares.

Como funciona

A ideia central é não construir uma infraestrutura de verificação separada e cara, mas se conectar a algo que já funciona. O autotuner de desempenho já testa todos os dias diferentes variantes de compilação do mesmo kernel em busca da mais rápida. Os autores adicionaram a esse ciclo uma comparação de resultados: a saída de cada variante é conferida com a do kernel fundido padrão. Se há divergência, alguma das variantes está sendo compilada incorretamente.

Se a suspeita recai sobre um bug dentro do grafo, o sistema também compara o resultado com uma referência não fundida, em que as operações são executadas separadamente. Assim é possível saber se a culpa é justamente da fusão.

Em seguida entra um isolador de dois estágios. Simplesmente dividir o kernel em partes costuma mascarar o bug, porque o problema vive justamente na versão fundida. Por isso, o isolador lê passo a passo os registradores vetoriais de hardware dentro do próprio kernel fundido e encontra a instrução em que os valores divergem pela primeira vez. Além disso, os autores usam geração de programas sintéticos e um enxame de agentes de IA que escrevem patches para os problemas encontrados.

O que foi encontrado

Nos clusters de TPU do Google, o sistema ajudou a corrigir mais de 142 bugs que o CI deixou passar. A distribuição é curiosa: 43% deles não estavam nas transformações de alto nível do grafo, mas nos geradores de código de baixo nível para registradores vetoriais e memória. Ou seja, quebrava onde os testes comuns em nível de grafo quase não olham.

Limitações e conclusões

Para ser justo, o trabalho tem pontos fracos. Todos os 142 bugs são contados em um único número, sem ablações por componente, então não está claro qual parte do sistema traz mais benefício: o fuzzing, o isolador ou os agentes. Condições de corrida de memória não determinísticas, os verdadeiros heisenbugs, ficam fora do alcance do método. Além disso, comparar com o kernel padrão pressupõe que ele próprio esteja correto, o que nem sempre acontece.

A conclusão prática vai além da TPU. Se você tem um autotuner, um compilador JIT ou qualquer mecanismo que percorra implementações equivalentes de uma mesma operação, já está pagando para calcular várias respostas para as mesmas entradas. Adicionar uma conferência de resultados com tolerância a ruído numérico é barato, e ela detecta justamente os erros que ninguém veria. Para equipes que treinam modelos grandes em sua própria stack, essa é uma das formas mais baratas de pegar defeitos antes que consumam o orçamento.