Проблема, которую видит каждый, но решить сложнее, чем кажется
Представьте, что вы хотите обучить большую языковую модель – скажем, аналог тех, что умеют писать тексты, отвечать на вопросы и вести диалог. Допустим, модель содержит 7 миллиардов параметров. Звучит внушительно, но давайте переведём это в практику: сколько памяти нужно, чтобы всё это тренировать?
Каждый параметр модели – это число с плавающей точкой. И во время обучения нейросеть помнит не только само это число, но и несколько его «спутников»: градиент (то, в какую сторону нужно изменить параметр) и переменные состояния оптимизатора (что-то вроде памяти о предыдущих шагах обучения, чтобы двигаться в нужном направлении точнее). Каждый из этих «спутников» занимает 4 байта. В сумме один параметр может тянуть за собой до 16 байтов памяти.
Умножьте 16 байтов на 7 миллиардов параметров – и получите больше 100 гигабайтов памяти ускорителя. Это больше, чем есть у большинства исследователей. Современные профессиональные GPU, вроде тех, что используются в научных лабораториях, нередко имеют 40 или 80 гигабайтов памяти – и даже этого не хватает. Обучение крупных моделей превращается в задачу, доступную только организациям с дорогостоящей инфраструктурой.
Именно здесь появляется FlashOptim – набор техник, разработанных для того, чтобы радикально сократить потребление памяти в процессе обучения, не потеряв при этом качество модели.
Что такое «смешанная точность» и почему вокруг неё столько всего
Перед тем как разбираться с FlashOptim, нужно понять один базовый принцип – смешанная точность обучения. Это звучит сложно, но идея простая.
Числа в компьютере можно хранить с разной точностью. Например, 32-битное число – это что-то вроде записной книжки, где можно зафиксировать очень много знаков после запятой. 16-битное – более краткая запись, менее точная, зато занимающая вдвое меньше места. При обучении нейросетей 32-битная точность нужна для одних операций (например, для обновления весов оптимизатором), а 16-битная достаточна для других (например, для самих вычислений в сети).
Смешанная точность – это компромисс: где можно, используем 16-битные числа, а где критически важна точность – 32-битные. Это уже экономит память, но недостаточно. Оптимизатор всё равно хранит свои внутренние переменные в 32 битах, и именно они «съедают» большую часть памяти.
FlashOptim идёт дальше и предлагает две ключевые техники: умное разбиение основных весов и сжатие переменных состояния оптимизатора.
Техника первая: умное разбиение весов
Во время стандартного обучения со смешанной точностью существует так называемая «мастер-копия» весов – 32-битная версия каждого параметра, которую оптимизатор обновляет на каждом шаге. Эта копия нужна, потому что при 16-битном хранении параметров маленькие изменения могут просто «потеряться» из-за недостаточной точности. Мастер-копия страхует от этих потерь.
Но хранить полную 32-битную копию для каждого параметра – расточительство. Авторы FlashOptim предлагают другой подход, который называется Master Weight Splitting (разбиение основных весов). Идея в том, чтобы не хранить всю 32-битную копию, а хранить только разницу между ней и 16-битной версией – так называемую дельту.
Представьте, что у вас есть точная карта города в высоком разрешении (32 бита) и упрощённая туристическая карта (16 бит). Вместо того чтобы носить с собой оба варианта, вы берёте туристическую карту и маленький листок с пометками: «здесь улица чуть-чуть шире, там здание снесли». Этот листок с пометками – и есть дельта.
Авторы FlashOptim пошли дальше классического подхода: они доказали, что эту дельту можно хранить не в 16 битах, а всего в 8 битах – при условии, что ошибка квантования (то есть потеря точности при переходе к меньшему числу битов) не превышает определённый порог. Они вывели точную математическую границу этой ошибки и показали: если обновление параметра укладывается в допустимый диапазон, 8-битная дельта справляется отлично. Если обновление выходит за пределы – систему можно аккуратно масштабировать или в редких случаях временно переключиться на 16-битный формат.
В результате параметры хранятся как 16-битная копия плюс 8-битная дельта – итого 3 байта вместо 4. Экономия кажется небольшой, но при миллиардах параметров цифры становятся впечатляющими.
Техника вторая: сжатие памяти оптимизатора через компандирование
Теперь самое интересное – переменные состояния оптимизатора. Для понимания нужно немного рассказать о том, зачем они вообще нужны.
Популярный оптимизатор AdamW, который используется при обучении большинства современных языковых моделей, отслеживает два «момента» для каждого параметра. Первый момент – это среднее значение градиентов за последние шаги (грубо говоря, в каком направлении двигались последнее время). Второй момент – это среднее значение квадратов градиентов (насколько сильно менялось это направление). Зная оба числа, AdamW может делать адаптивные шаги: двигаться быстро там, где уверен, и осторожно – там, где градиенты «прыгают».
Представьте опытного скалолаза, который помнит не только куда лезет, но и насколько ненадёжен камень в каждой точке. Это и есть смысл двух моментов. Каждый из них обычно хранится в 32-битном формате – итого 8 байтов только на состояние оптимизатора.
FlashOptim предлагает хранить эти моменты в 8-битном формате. Но здесь возникает проблема: прямолинейное сжатие 32-битного числа до 8-битного сильно снижает точность. Особенно плохо это работает для значений, которые сильно разбросаны по шкале – а значения моментов именно такие: большинство маленькие, но встречаются и крупные.
Для решения этой проблемы авторы используют компандирование – слово из телефонии, образованное от «компрессия» и «расширение». Идея в том, чтобы перед сохранением нелинейно «сжать» шкалу значений, а при чтении – «растянуть» обратно. Это позволяет использовать больше 8-битных «слотов» для маленьких значений (где нужна точность) и меньше – для больших (где точность менее критична).
Хорошая аналогия – звуковая запись. Тихие звуки в записи важно передать точно, потому что именно они создают нюансы музыки. Громкие звуки более «грубые» и прощают большую погрешность. Именно поэтому в аудиоформатах шкала громкости нелинейная – тихие уровни размечены гуще, громкие – реже. Компандирование делает то же самое для чисел в оптимизаторе.
Авторы FlashOptim разработали два варианта компандирования:
- Аналитическое логарифмическое компандирование. Функция сжатия строится на основе логарифмической шкалы – математически обоснованный выбор, который хорошо работает для типичных распределений значений моментов. Его можно задать заранее, и он не требует дополнительного обучения.
- Обучаемое компандирование. Здесь параметры функции сжатия сами обучаются в процессе тренировки модели. Небольшая вспомогательная сеть адаптируется к конкретным распределениям значений для конкретной задачи – это потенциально точнее, но чуть сложнее в реализации.
Оба подхода позволяют хранить каждую переменную состояния оптимизатора в 1 байте вместо 4. Для AdamW с двумя моментами это сокращение с 8 байтов до 2.
Третий компонент: 16-битные градиенты
Здесь авторы FlashOptim не изобретают ничего нового – они просто сохраняют уже сложившуюся практику. В стандартном обучении со смешанной точностью градиенты уже давно хранятся в 16-битном формате (форматы FP16 или BF16). Этого достаточно для точного обновления весов, и снижения качества не наблюдается. Так что 2 байта на параметр для градиентов остаются без изменений.
Дополнительно, если использовать стратегию немедленного освобождения градиентов (то есть удалять их сразу после обновления весов, не ожидая конца шага), память под градиенты можно вообще не учитывать постоянно – и тогда суммарный расход падает до 5 байтов на параметр.
Итог в байтах: было и стало
Теперь соберём всё воедино. При стандартном обучении со смешанной точностью для оптимизатора AdamW один параметр требует:
- 4 байта – 32-битная мастер-копия параметра,
- 2 байта – 16-битный рабочий параметр (используется в вычислениях),
- 2 байта – градиент в формате BF16,
- 4 байта – первый момент оптимизатора,
- 4 байта – второй момент оптимизатора.
Итого – около 16 байтов на параметр. С FlashOptim картина меняется:
- 2 байта – 16-битная копия параметра,
- 1 байт – 8-битная дельта,
- 2 байта – градиент в формате BF16,
- 1 байт – первый момент в 8-битном квантованном формате,
- 1 байт – второй момент в 8-битном квантованном формате.
Итого – 7 байтов. А с освобождением градиентов – 5 байтов. Это более чем двукратное сокращение потребления памяти по сравнению с исходной цифрой в 16 байтов.
Для модели с 7 миллиардами параметров это означает переход от требования в 100+ гигабайт к примерно 35–50 гигабайтам. Разница между «нужен специальный вычислительный кластер» и «можно запустить на одном мощном GPU».
Что показали эксперименты
Авторы проверяли FlashOptim на нескольких оптимизаторах: SGD, AdamW и Lion. Задачи включали компьютерное зрение (классификация изображений на датасете ImageNet с использованием архитектур ResNet и Vision Transformers) и задачи обработки естественного языка – в том числе тонкую настройку модели Llama-3.1-8B.
Результат, который авторы подчёркивают особо: ни на одном из тестов измеримого ухудшения качества не зафиксировано. Точность классификации изображений осталась в пределах статистической погрешности. Языковые модели после тонкой настройки с FlashOptim показали сопоставимые результаты по метрике перплексии (это показатель того, насколько уверенно модель предсказывает текст – чем ниже, тем лучше) и другим задачам понимания языка.
Это критически важный результат. Методы сжатия нередко дают выигрыш в памяти ценой падения точности модели – и именно этот компромисс обычно останавливает их внедрение. FlashOptim, судя по опубликованным экспериментам, позволяет обойти этот компромисс.
Ещё один бонус – сокращение размера контрольных точек (checkpoint). Контрольная точка – это сохранённый «слепок» состояния модели в определённый момент обучения: параметры, состояния оптимизатора, всё вместе. Поскольку FlashOptim хранит состояния компактнее, контрольные точки уменьшаются более чем вдвое. При работе с большими моделями это означает существенную экономию дискового пространства и значительно более быструю передачу файлов.
Почему это важно для тех, кто не работает в BigTech
Обучение крупных нейронных сетей долгое время было уделом организаций, располагающих сотнями GPU и петабайтами хранилищ. Небольшие исследовательские группы, университеты и независимые лаборатории оказывались фактически отрезаны от работы с такими моделями – не из-за отсутствия идей или алгоритмов, а просто из-за памяти.
FlashOptim меняет это уравнение. Он работает как надстройка над существующими оптимизаторами – SGD, AdamW, Lion – и совместим с популярными фреймворками глубокого обучения, такими как PyTorch. Это означает, что для его применения не нужно переписывать весь код обучения. Достаточно заменить оптимизатор – и экономия памяти достигается почти автоматически.
Совместимость с API (то есть с программным интерфейсом, через который разработчики взаимодействуют с библиотеками) – это не технический каприз, а принципиальный выбор. Именно она определяет, будет ли инструмент реально использоваться или останется красивым результатом в научной статье.
Насколько сложно то, что происходит внутри
За простотой использования стоит нетривиальная математика. Квантование – сжатие числа из одного формата в другой с меньшим числом битов – всегда вносит ошибку. Вопрос в том, насколько эта ошибка критична для итогового качества модели.
Для весов и градиентов ошибка квантования может быть катастрофической – именно поэтому параметры модели обычно хранятся в высокоточных форматах. Для состояний оптимизатора – первого и второго моментов AdamW – ситуация иная: эти числа менее чувствительны к точности, и 8-битного представления при правильно подобранной функции компандирования оказывается достаточно.
Авторы FlashOptim провели детальный анализ: они формально доказали, что ошибка квантования при 8-битном хранении дельты весов остаётся ниже порога, при котором она влияет на обучение. Это не просто эмпирическое наблюдение («попробовали – работает»), а математически обоснованное утверждение – что делает подход значительно более надёжным.
Аналогия из повседневной жизни: если вы переводите длинное число типа «3,14159265358979» в «3,14», для большинства практических расчётов это вполне приемлемо. Но если число используется в навигации космического аппарата – каждая цифра на счету. FlashOptim методично выясняет, в каких именно «вычислениях» нейросети достаточно «3,14» – и использует это знание для экономии памяти там, где это безопасно.
Два режима компандирования: выбор за задачей
Интересная деталь – FlashOptim не навязывает единственный способ сжатия состояний оптимизатора. Два варианта компандирования ориентированы на разные ситуации:
Аналитическое логарифмическое компандирование подходит для случаев, когда нужна предсказуемость и простота. Функция заранее определена, не требует дополнительных вычислений в процессе обучения и хорошо работает для типичных распределений значений моментов. Это «базовая комплектация».
Обучаемое компандирование – это «расширенная версия». Параметры функции сжатия подстраиваются под конкретную модель и конкретную задачу прямо в процессе обучения. Это требует чуть больших накладных расходов, зато потенциально даёт более точное квантование для нестандартных распределений. Практические эксперименты показали, что оба метода работают хорошо – что само по себе говорит о надёжности общей концепции.
Что это означает на практике
FlashOptim – это не революция в архитектуре нейросетей и не прорыв в теории оптимизации. Это точная инженерная работа: найти, где в процессе обучения тратится память впустую, доказать, что именно там можно сэкономить без вреда для результата, и реализовать это аккуратно и совместимо с существующим кодом.
Результат – снижение требований к памяти более чем вдвое при сохранении качества модели, совместимость с популярными оптимизаторами и фреймворками, сокращение размера контрольных точек и возможность обучать модели масштаба 7 миллиардов параметров на оборудовании с 35–50 гигабайтами памяти вместо 100+.
В мире, где доступ к вычислительным ресурсам всё больше определяет, кто может вести исследования на переднем крае, такие инструменты – как FlashOptim – смещают это неравенство в сторону большей доступности. И это, пожалуй, не менее важный результат, чем сама техника.
