NVIDIA показала побитовый детерминизм в Megatron Core для претрейна на тысячах GPU

Обучение модели с триллионами параметров на тысячах GPU стоит миллионы долларов и занимает недели. Если в середине прогона всплывает всплеск лосса или NaN, повторить его обычно не удаётся: порядок операций меняется от запуска к запуску, и ошибка не проявляется. NVIDIA опубликовала на Developer Blog материал о том, как добиться побитового детерминизма при претрейне в Megatron Core. Два прогона с одинаковыми данными и конфигурацией должны давать идентичные результаты до последнего бита.

Почему детерминизм в претрейне такая редкость

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

В источнике перечислены усложняющие факторы:

• несколько измерений параллелизма одновременно (данные, тензоры, конвейер и другие), каждое со своими редукциями и обменами;
• вычисления в низкой точности, где ошибки округления крупнее и сильнее зависят от порядка операций;
• распределённые чекпойнты, которые нужно сохранять и загружать так, чтобы продолжение обучения не меняло траекторию;
• масштаб: триллионы параметров и тысячи GPU, где вероятность редкого сбоя за прогон близка к единице.

Конкретные приёмы, которыми Megatron Core обеспечивает порядок операций, и стоимость детерминизма по скорости в доступном фрагменте не раскрыты. Эти данные нужно смотреть в полной версии материала.

Что даёт побитовая воспроизводимость

В анонсе названы три практические выгоды: отладка, валидация и воспроизводимое возобновление обучения.

Первая касается поиска причин сбоя. Если прогон детерминирован, можно перезапустить его с последнего чекпойнта перед аномалией и получить ту же аномалию. Это превращает «призрачную» ошибку в обычный баг, на который можно ставить проверки и бисекцию по слоям или шагам.

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

Третья касается рестартов. На больших кластерах отказы оборудования случаются регулярно. Если после загрузки чекпойнта обучение идёт ровно по той же траектории, что и без остановки, инженеры могут быть уверены, что перезапуск не внёс скрытого дрейфа.

Цена вопроса и где это пригодится

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

При этом логика применения понятна. Постоянно держать детерминированный режим в самых дорогих прогонах не обязательно. Но он полезен в нескольких сценариях:

• при воспроизведении и разборе сбоев, когда нужен точный повтор;
• при проверке новых ядер, смены точности (например, перехода на более низкую) или изменения схемы параллелизма;
• при регрессионном тестировании кода обучения: эталонный прогон на небольшом числе шагов, с которым побитово сравнивается каждая сборка.

Практический вывод для команд

Командам, которые обучают собственные модели на Megatron Core, стоит завести короткий контрольный прогон в детерминированном режиме и включить его в CI. Любое изменение в коде, драйверах или конфигурации параллелизма сначала сравнивается с эталоном по хешам весов после фиксированного числа шагов. Расхождение сразу показывает, что изменение затронуло численное поведение.

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