Google встроил фаззинг в автотюнер XLA и поймал 142 тихих бага компилятора на TPU

Если компилятор падает с ошибкой, вы это заметите. Если он молча выдаёт неверные числа, вы узнаете об этом через неделю обучения, когда лосс поведёт себя странно, а счёт за кластер уже придёт. Именно такие тихие баги команда Google разбирает в работе XLA from the Trenches, и решение у них неожиданно приземлённое.

Почему тихие баги хуже падений

XLA превращает граф вычислений модели в машинный код для TPU. По пути он сливает операции в большие ядра (fusion), раскладывает данные по векторным регистрам и памяти и делает ещё десятки оптимизаций. Каждая из них может быть корректной по отдельности и ломаться в редком сочетании форм, типов и параметров.

Статические тесты в CI проверяют заранее придуманные случаи. Миллионы таких тестов не покрывают все сочетания, которые рождаются в реальных моделях. Баг при этом не проявляется как падение: всё запускается, всё быстро, просто результат чуть-чуть не тот. На обучении большой модели цена такой ошибки измеряется миллионами долларов.

Как это устроено

Ключевая идея в том, чтобы не строить отдельную дорогую инфраструктуру проверки, а подключиться к тому, что уже работает. Автотюнер производительности и так каждый день перебирает разные варианты компиляции одного и того же ядра в поисках самого быстрого. Авторы добавили в этот цикл сравнение результатов: ответ каждого варианта сверяется с ответом дефолтного слитого ядра. Если расхождение есть, один из вариантов компилируется неправильно.

Если подозрение падает на баг внутри графа, система сравнивает результат ещё и с неслитым эталоном, где операции выполняются по отдельности. Так можно понять, виноват ли именно fusion.

Дальше подключается двухэтапный изолятор. Простое разбиение ядра на части часто маскирует баг, потому что проблема живёт именно в слитой версии. Поэтому изолятор пошагово читает аппаратные векторные регистры внутри самого слитого ядра и находит инструкцию, где значения впервые расходятся. Кроме того, авторы используют генерацию синтетических программ и рой ИИ-агентов, которые пишут патчи для найденных проблем.

Что нашли

На кластерах TPU в Google система помогла устранить более 142 багов, пропущенных в CI. Любопытно распределение: 43% из них сидели не в высокоуровневых преобразованиях графа, а в низкоуровневых генераторах кода для векторных регистров и памяти. То есть ломалось там, куда обычные тесты на уровне графа почти не заглядывают.

Ограничения и что делать

Честно говоря, у работы есть слабые места. Все 142 бага посчитаны одной цифрой, без абляций по компонентам, поэтому непонятно, какая часть системы приносит больше всего пользы: фаззинг, изолятор или агенты. Недетерминированные гонки памяти, то есть настоящие гейзенбаги, остаются за рамками метода. Кроме того, сравнение с дефолтным ядром предполагает, что дефолт сам правильный, а это не всегда так.

Практический вывод шире, чем TPU. Если у вас есть автотюнер, JIT-компилятор или любой механизм, который перебирает эквивалентные реализации одной операции, вы уже платите за вычисление нескольких ответов на одни и те же входы. Добавить сверку результатов с допуском на численный шум стоит дёшево, а ловит она именно те ошибки, которые никто не увидит. Для команд, которые обучают большие модели на своём стеке, это один из самых недорогих способов поймать брак до того, как он съест бюджет.