Sakana AI совместно с NVIDIA: позволяет GPU пропускать до 80% бесполезных вычислений больших моделей, ускорение вывода H100 на 30%

По данным мониторинга Beating, Sakana AI совместно с NVIDIA открыли формат разреженных данных под названием TwELL и сопутствующий ускоряющий ядро, успешно позволяющее GPU при работе с крупными моделями пропускать «бесполезные вычисления», результаты которых близки к нулю. Эта схема при сохранении точности модели позволяет увеличить скорость вывода H100 до 30%, ускорить обучение до 24% и значительно снизить пиковое использование видеопамяти.

Передний слой больших моделей (FFN) занимает большую часть параметров и вычислительных ресурсов. Но на самом деле при каждом генерации текста более 80% нейронов находятся в «спящем режиме» (активность близка к нулю), не внося никакого вклада в конечный результат. Если удастся пропускать эти нейроны, можно сэкономить огромное количество вычислительной мощности. Однако современные GPU изначально оптимизированы для выполнения плотных матриц, и при использовании традиционных методов для выделения разбросанных полезных данных затраты на поиск и чтение данных могут полностью съесть выгоду от экономии вычислений.

Формат TwELL специально создан, чтобы разрушить этот аппаратный барьер. Он полностью ориентирован на параллельную архитектуру GPU: вместо того чтобы объединять ненулевые данные по регионам, как в традиционных методах, он разбивает данные на небольшие блоки (тайлы), которые GPU обрабатывает наиболее эффективно. Таким образом, вычислительные ядра GPU могут напрямую собирать полезные данные локально, полностью устраняя задержки на глобальное чтение и запись видеопамяти, что идеально вписывается в ускорительный конвейер современных чипов.

На практике при обучении модели с 1,5 миллиарда параметров достаточно было ввести небольшую регуляризацию, чтобы снизить долю нейронов, требующих вычислений, до менее 2%, при этом производительность по семи задачам не ухудшилась. Данные также выявили закономерность: чем больше параметров у модели, тем больше нейронов находится в спящем режиме (например, у модели с 2 миллиардами параметров доля ненулевых элементов на 38% ниже, чем у модели с 500 миллионами). Это означает, что в будущем при создании еще более крупных моделей эта оптимизация, ориентированная на низкоуровневое аппаратное обеспечение, сможет дать еще более значительный прирост производительности.