Оптимизация больших языковых моделей с помощью Puzzle: новый подход к ускорению вывода
В последние годы большие языковые модели (LLM) продемонстрировали выдающиеся способности в различных задачах, от генерации текста до обработки естественного языка. Однако их широкое применение ограничено высокими вычислительными затратами, особенно во время вывода. Исследование, представленное в работе "Puzzle: Distillation-based NAS for Inference-optimized LLMs", предлагает инновационный подход к этой проблеме, сочетая методы поиска архитектуры нейронных сетей (NAS) и дистилляцию знаний для оптимизации LLM.
С увеличением размера языковых моделей наблюдается рост их точности и возможностей. Тем не менее, это также ведет к увеличению вычислительных затрат на вывод, что делает их менее доступными для широкого применения. Работа "Puzzle" предлагает решение этой проблемы, позволяя оптимизировать модели для конкретного аппаратного обеспечения, сохраняя при этом их основные способности.
Проблема избыточности
Одной из основных проблем, с которой сталкиваются LLM, является избыточность параметров. После завершения обучения многие из параметров оказываются избыточными для задач вывода. Это создает возможность для оптимизации моделей с целью снижения вычислительных затрат без значительной потери точности. Однако традиционные архитектуры LLM в основном остаются однородными, что не учитывает различия в вычислительных затратах и вкладе каждого блока в общую производительность модели.
Подход Puzzle
Структура Puzzle
Puzzle представляет собой многоступенчатую архитектуру, которая включает в себя три основных этапа:
-
Создание библиотеки блоков: На этом этапе проводятся эксперименты с различными вариантами блоков (например, слоев внимания и полносвязных слоев), которые обучаются параллельно с использованием методики блочной локальной дистилляции (BLD).
-
Оценка архитектур: Каждый блок оценивается с точки зрения его качества, и выбираются наиболее эффективные конфигурации для дальнейшей сборки модели.
-
Глобальная дистилляция знаний (GKD): На последнем этапе модель дообучается для улучшения ее производительности на основе знаний, полученных от родительской модели.
Блочная локальная дистилляция (BLD)
BLD позволяет обучать подмодули (блоки) независимо друг от друга, что значительно снижает вычислительные затраты. Каждый блок обучается с использованием активаций родительского блока, что способствует более быстрой сходимости и позволяет использовать меньшие объемы данных для обучения.
Поиск архитектуры
Для поиска оптимальных конфигураций используется смешанное целочисленное программирование (MIP), что позволяет эффективно исследовать огромное пространство архитектур, состоящее из различных вариантов блоков. Это позволяет находить наилучшие архитектуры, соответствующие заданным ограничениям по памяти, задержкам и пропускной способности.
Результаты
Работа демонстрирует эффективность подхода Puzzle через создание модели Nemotron-51B, производной от Llama-3.1-70B-Instruct. Nemotron-51B достигает 2,17-кратного увеличения пропускной способности вывода, что позволяет запускать модель на одном графическом процессоре NVIDIA H100, сохраняя 98,4% точности родительской модели. Для обучения этой модели потребовалось всего 45 миллиардов токенов, что значительно меньше, чем 15 триллионов токенов, использованных для обучения родительской модели.
Сравнение с традиционными подходами
Традиционные методы, такие как обучение с нуля или дистилляция знаний, часто требуют значительных вычислительных ресурсов и времени. Puzzle позволяет значительно сократить эти затраты, обеспечивая при этом высокую производительность моделей. Это делает его особенно полезным для практиков, которые хотят адаптировать существующие модели под конкретные задачи и аппаратные условия.
Заключение
Работа "Puzzle" представляет собой важный шаг в оптимизации больших языковых моделей для вывода. Используя методы NAS и дистилляции знаний, исследование демонстрирует, что можно значительно улучшить эффективность LLM, сохраняя при этом их способности. Это открывает новые горизонты для применения мощных языковых моделей в реальных сценариях, делая их более доступными и эффективными для широкого круга пользователей.
Таким образом, Puzzle не только предоставляет инструменты для оптимизации существующих моделей, но и закладывает основы для будущих исследований в области адаптации архитектур LLM к конкретным задачам и аппаратным условиям.