Представьте: вы обучаете языковую модель, всё настроено, система запущена – и вдруг начинают появляться странные числа. Не ошибки, не вылеты, а просто тихое расхождение в вероятностях, которое нельзя объяснить ничем очевидным. Именно с этого начался один из самых поучительных эпизодов в практике команды AI21 – компании, разрабатывающей собственные языковые модели.
«Одно и то же, но другое»
При обучении языковых моделей с использованием метода GRPO – одного из подходов к обучению с подкреплением – есть важная проверка. Система генерирует текст, записывает, с какой «уверенностью» она выбирала каждое слово, а затем та же самая модель (с теми же весами, без каких-либо обновлений) пересчитывает эти значения заново. Результаты должны совпадать практически идеально: те же веса, те же входные данные, тот же результат.
Но в случае с Jamba 3B – гибридной моделью AI21, сочетающей стандартные механизмы внимания с архитектурой Mamba – совпадения не было. Числа расходились. Причём расходились не хаотично, а с определённой периодичностью: сбой появлялся примерно каждые 12 шагов обучения, потом «пропадал», а затем возвращался снова.
Самое неприятное – внешне это выглядело как обычная нестабильность обучения. Подобные вещи случаются, и списать всё на «шум» было бы очень легко.
Поиск рычага
Распределённые системы обучения – сложные конструкции. В них одновременно работают движок для генерации текста, система для обновления весов, координация между несколькими машинами и передача данных между компонентами. Когда что-то ломается в такой системе, велик соблазн начать отлаживать всё сразу и зайти в тупик.
Команда выбрала другой путь: найти параметр, изменение которого меняет характер сбоя, а не просто его интенсивность. Проще говоря – найти рычаг.
Им оказалось количество генерируемых текстов на каждый запрос. Когда исследователи начали увеличивать это число – 8, 16, 32, 64, 128 – они заметили важное: периодичность сбоев менялась вместе с этим параметром. При 128 текстах на запрос сбой появлялся уже на самом первом шаге.
Это наблюдение изменило всю картину. Если сбой синхронизирован с процессом генерации, значит, проблема, скорее всего, именно там – а не в системе обновления весов, не в синхронизации между машинами и не в самом алгоритме обучения.
От «странности при обучении» к «воспроизводимому дефекту»
Следующая цель – воспроизвести баг в как можно более простых условиях. В идеале – на нулевом шаге, до того как обучение вообще началось. Это важно: когда ошибка появляется с первой же итерации, из уравнения выпадают все накопленные эффекты – история градиентов, дрейф весов и долгосрочная динамика обучения. Остаётся просто выполнение кода.
Со 128 текстами на запрос это удалось. Баг воспроизводился на первом же шаге, надёжно и стабильно.
Дальнейшее сужение круга подозреваемых дал ещё один параметр – объём GPU-памяти, выделяемой под кэш. При значении 50% баг исчезал. При большем значении – появлялся снова. Это означало, что дело в том, как движок инференса (то есть генерации текста) выделяет и использует кэш внутри видеокарты.
А поскольку Jamba – гибридная архитектура, кэш у неё бывает двух типов: один для механизма внимания, другой для блоков Mamba. Проверили модель только с механизмом внимания – баг не воспроизвёлся. Значит, дело в Mamba-части.
Два символа и несколько недель
Источник проблемы нашли в низкоуровневом коде, который выполняется непосредственно на GPU. Там есть операция: вычислить, куда именно в памяти записать состояние для каждого элемента кэша. Для этого нужно перемножить два числа: номер элемента и размер одного «шага» в памяти.
Оба числа были объявлены как 32-битные беззнаковые целые. Это тип, который умеет хранить значения примерно до 4,29 миллиарда. Звучит внушительно, но произведение этих двух чисел могло запросто выйти за этот предел.
Проще говоря: представьте одометр, который показывает только до 999 999 км, а потом сбрасывается в ноль. Машина проехала миллион километров – одометр показывает 000 000. Никакой ошибки, никакого предупреждения. Просто неправильное число.
Именно это и происходило. При размере шага 89 600 элементов переполнение наступало, когда номер элемента превышал примерно 47 935. В реальной конфигурации кэш содержал 69 776 слотов – то есть около 31% из них записывались по неверным адресам в памяти GPU. Данные уходили «не туда», а в «правильных» местах оставались нули.
Никакого сбоя системы. Никакого предупреждения. Просто тихо идущие неправильные числа на выходе, которые затем интерпретировались как нестабильность при обучении.
Исправление заняло две секунды: тип данных заменили с 32-битного на 64-битный. Теперь числа не переполняются. Всё.
Несколько недель расследования. Два изменённых символа в коде.
Почему это сложно поймать
Подобные баги особенно коварны в системах машинного обучения по нескольким причинам.
Во-первых, они не ломают систему явно. Нет исключения, нет неверного формата данных, нет очевидного сигнала. Есть просто чуть другие числа, которые вполне могли бы быть следствием сотни других причин.
Во-вторых, они проявляются только при определённом масштабе. При небольшом числе запросов кэш не заполнялся до критической отметки и переполнения не происходило. Нужно было специально «надавить» на систему, чтобы баг вообще себя проявил.
В-третьих, в распределённых системах обучения симптом и причина могут находиться очень далеко друг от друга. Ошибка в GPU-ядре на этапе генерации текста выглядела как нестабильность на этапе обучения – в совершенно другом компоненте системы.
Урок не про баг, а про подход
Самое ценное в этой истории – не сам дефект, а метод его обнаружения.
Когда система ведёт себя странно, первый инстинкт – проверить всё сразу. Это редко помогает. Гораздо эффективнее искать параметр, который меняет структуру сбоя. Не «при каком значении ошибка становится больше», а «при каком значении она начинает вести себя по-другому». Именно это даёт подсказку о природе проблемы.
Второй принцип – сужать до минимума. Если баг можно воспроизвести на первом шаге вместо пятисотого, нужно добиться именно этого. Чем меньше контекст, тем чище сигнал.
Третий – изолировать подсистемы. Большой распределённый конвейер – плохое место для отладки. Хорошее место – минимальный скрипт, который воспроизводит проблему в изоляции.
Эти принципы работают не только для GPU-ядер и не только в машинном обучении. Это просто хорошая инженерная практика: когда не понимаешь, что сломалось, сначала пойми, где сломалось. А для этого нужно уметь задавать системе вопросы и правильно интерпретировать ответы.
