Thinking Machines ускорила обучение Inkling 2 Small 4-битными числами

Лаборатория Миры Мурати Thinking Machines ускорила обучение своей модели Inkling 2 Small: часть вычислений на всех этапах, от предобучения¹ до обучения с подкреплением², шла в 4-битных числах. По словам компании, это крупнейшая известная ей модель, которую так обучили от начала до конца.

Почему это важно: чем меньше битов уходит на число, тем меньше нужно памяти и тем быстрее считает видеокарта, но тем грубее приходится округлять. Ошибки округления копятся от слоя к слою и от шага к шагу и могут сорвать обучение. Thinking Machines выиграла в скорости, почти не потеряв в качестве. К тому же сжатую 4-битную версию готовой модели не нужно отдельно подстраивать (калибровать), как обычно бывает при сжатии.

Предыстория: Inkling 2 Small компания выпустила 5 ноября. Это смесь экспертов³ на 594 млрд параметров, из которых при обработке каждого токена работают 17 млрд.

В цифрах: выигрыш и цена 4-битного счёта, по замерам компании.

• На видеокартах NVIDIA GB300 формат NVFP4 занимает в 3,6 раза меньше битов, чем привычный 16-битный BF16, а считается до 6 раз быстрее.
• По сравнению с обучением целиком в 8-битном формате MXFP8 в самых крупных запусках предобучение ускорилось на 8%, обучение с подкреплением — на 20%.
• Плата — около 2% эффективности при предобучении: в 8 битах модель достигла бы той же ошибки, потратив примерно на 2% меньше вычислений. В обучении с подкреплением разницы в награде компания не заметила.
• Предобучение Inkling 2 Small заняло 68 часов, а несколько месяцев назад модель похожего размера училась несколько недель. Но это итог многих доработок, а не одного нового формата.

Детали: четыре бита вмещают всего 16 значений — от −6 до 6, например 0,5, 1,5, 3 или 4, но не 5. Чтобы записывать числа крупнее, NVFP4 делит данные на блоки по 16 чисел и хранит для каждого общий множитель: если самое большое по модулю число в блоке равно −54, оно станет −6 с множителем 9. Обучать модель целиком в 4 битах не вышло: первые опыты ударили по качеству сильнее, чем компания готова была допустить. Поэтому NVFP4 включён только при прямом проходе, когда модель вычисляет ответ, и только в слоях экспертов: на них уходит много вычислений, а точность от сжатия там страдает меньше, чем в других слоях.

Как это работает: обратный проход, где по ошибке модели вычисляются поправки к её весам, идёт в 8-битном MXFP8. Он чувствителен к сдвигу, когда числа раз за разом округляются в одну сторону. Таблицы чисел из прямого прохода здесь приходится переворачивать: блоки по 16 чисел, собранные вдоль строк, теперь идут вдоль столбцов. Каждое число оказывается рядом с 15 новыми соседями, получает другой множитель и округляется иначе, поэтому поправки к весам могут выйти неточными. Для весов NVIDIA предлагает один общий множитель на квадрат 16 на 16 чисел: при перевороте он не меняется. Но в квадрате из 256 чисел одно крупное значение может поднять множитель так, что остальные округлятся до нуля, и для компании такой способ оказался слишком грубым. Thinking Machines переводит числа в более точный формат и округляет заново случайным образом: например, 2,25 станет двойкой в трёх случаях из четырёх и тройкой в одном. В среднем выходит ровно 2,25, и ошибки не копятся в одну сторону.

Но: ускорив расчёты, компания столкнулась с новым узким местом — пересылкой данных между видеокартами. Эксперты модели разложены по разным видеокартам, и токены отправляются к нужным экспертам и обратно. Пока по соединению NVLink приходит один байт, видеокарта успевает выполнить до 16 000 операций, поэтому без сжатия её вычислительные блоки простаивают в ожидании. Данные стали сжимать и в пути: результаты экспертов идут по 8,25 бита на число вместо 16, и измеримой потери точности при обучении компания не увидела.

Общая картина: ещё в 2020 году нейросети по умолчанию учили в 32-битном формате FP32, пока видеокарты NVIDIA Ampere не получили поддержку 16-битного BF16. Теперь в Thinking Machines BF16 называют «высокой точностью», а в обучении используют 8 и даже 4 бита.

Что дальше: компания пробует перевести на 4 бита и обратный проход, а также изучает, как изменятся её приёмы на новых системах NVIDIA Vera Rubin.