Ускорение обучения LLM: новый метод от MIT и NVIDIA

Рассуждающие большие языковые модели (LLM) созданы для решения сложных задач, разбивая их на серию более мелких шагов. Эти мощные модели особенно хорошо справляются с задачами, требующими глубоких рассуждений, например, продвинутое программирование и многоэтапное планирование.

Однако разработка рассуждающих моделей требует огромного количества вычислений и энергии из-за неэффективности процесса обучения. В то время как некоторые высокопроизводительные процессоры непрерывно обрабатывают сложные запросы, другие простаивают.

Исследователи из Массачусетского технологического института и других организаций нашли способ использовать это время простоя для эффективного ускорения обучения рассуждающих моделей.

Их новый метод автоматически обучает небольшую, быструю модель для предсказания результатов работы более крупной рассуждающей LLM, которые затем проверяются большей моделью. Это снижает объем работы, которую должна выполнять рассуждающая модель, ускоряя процесс обучения.

Ключом к этой системе является ее способность обучать и развертывать небольшую модель адаптивным образом, чтобы она включалась только тогда, когда некоторые процессоры простаивают. Используя вычислительные ресурсы, которые в противном случае были бы потрачены впустую, она ускоряет обучение без дополнительных накладных расходов.

При тестировании на нескольких рассуждающих LLM метод удвоил скорость обучения, сохраняя при этом точность. Это может снизить стоимость и повысить энергоэффективность разработки продвинутых LLM для таких приложений, как прогнозирование финансовых трендов или выявление рисков в энергосетях.

"Люди хотят модели, которые могут справляться с более сложными задачами. Но если это цель разработки модели, то нам нужно уделять приоритетное внимание эффективности. Мы нашли без потерь решение этой проблемы и разработали полнофункциональную систему, которая может обеспечить значительное ускорение на практике", - говорит Цинхао Ху, постдок в MIT и соавтор работы, посвященной этой технике.

К нему присоединились соавтор Шан Янг, студент-аспирант факультета электротехники и компьютерных наук (EECS); Цзюньсянь Го, студент-аспирант факультета EECS; старший автор Сун Хан, доцент факультета EECS, член Исследовательской лаборатории электроники и выдающийся ученый NVIDIA, а также другие исследователи из NVIDIA, ETH Zurich, MIT-IBM Watson AI Lab и Университета Массачусетса в Амхерсте. Результаты исследования будут представлены на конференции ACM International Conference on Architectural Support for Programming Languages and Operating Systems.

Ограничение обучения

Разработчики хотят, чтобы рассуждающие LLM выявляли и исправляли ошибки в своем критическом мышлении. Эта возможность позволяет им успешно решать сложные запросы, которые могут оказаться не по силам обычной LLM.

Чтобы обучить их этому навыку, разработчики обучают рассуждающие LLM с использованием техники, называемой обучением с подкреплением (RL). Модель генерирует несколько возможных ответов на запрос, получает вознаграждение за лучший вариант и обновляется на основе лучшего ответа. Эти шаги повторяются тысячи раз по мере обучения модели.

Но исследователи обнаружили, что процесс генерации нескольких ответов, называемый rollout, может потреблять до 85 процентов времени выполнения, необходимого для обучения RL.

"Обновление модели — то есть фактическая "учебная" часть — потребляет очень мало времени по сравнению с этим", - говорит Ху.

Это ограничение возникает в стандартных алгоритмах RL, поскольку все процессоры в обучающей группе должны завершить свои ответы, прежде чем они смогут перейти к следующему шагу. Поскольку некоторые процессоры могут работать над очень длинными ответами, другие, которые сгенерировали более короткие ответы, ждут их завершения.

"Наша цель заключалась в том, чтобы превратить это время простоя в ускорение без каких-либо потерь", - добавляет Ху.

Они стремились использовать существующую технику, называемую спекулятивным декодированием, для ускорения процесса. Спекулятивное декодирование включает в себя обучение небольшой модели, называемой drafter, для быстрого предсказания будущих результатов работы большей модели.

Большая модель проверяет предположения drafter, и ответы, которые она принимает, используются для обучения.

Поскольку большая модель может одновременно проверять все предположения drafter, а не генерировать каждый результат последовательно, это ускоряет процесс.

Адаптивное решение

Но в спекулятивном декодировании модель drafter обычно обучается только один раз и остается статической. Это делает эту технику невозможной для обучения с подкреплением, поскольку рассуждающая модель обновляется тысячи раз в процессе обучения.

Статический drafter быстро устареет и станет бесполезным после нескольких шагов.

Чтобы преодолеть эту проблему, исследователи создали гибкую систему, известную как "Taming the Long Tail", или TLT.

Первая часть TLT — это адаптивный тренер drafter, который использует свободное время на простаивающих процессорах для обучения модели drafter на лету, поддерживая ее в хорошем соответствии с целевой моделью без использования дополнительных вычислительных ресурсов.

Второй компонент, адаптивный движок rollout, управляет спекулятивным декодированием для автоматического выбора оптимальной стратегии для каждой новой партии входных данных. Этот механизм изменяет конфигурацию спекулятивного декодирования на основе характеристик учебной нагрузки, таких как количество входных данных, обработанных моделью drafter, и количество входных данных, принятых целевой моделью во время проверки.

Кроме того, исследователи разработали модель drafter, чтобы она была легковесной, чтобы ее можно было быстро обучить. TLT повторно использует некоторые компоненты процесса обучения рассуждающей модели для обучения drafter, что приводит к дополнительным преимуществам в ускорении.

"Как только некоторые процессоры заканчивают свои короткие запросы и простаивают, мы немедленно переключаем их на обучение модели drafter, используя те же данные, которые они используют для процесса rollout. Ключевым механизмом является наше адаптивное спекулятивное декодирование — эти достижения были бы невозможны без него", - говорит Ху.

Они протестировали TLT на нескольких рассуждающих LLM, которые были обучены с использованием данных реального мира. Система ускорила обучение на 70–210 процентов, сохраняя при этом точность каждой модели.

В качестве дополнительного бонуса небольшую модель drafter можно легко использовать для эффективного развертывания в качестве бесплатного побочного продукта.

В будущем исследователи хотят интегрировать TLT в большее количество типов учебных и инференсных фреймворков и найти новые приложения обучения с подкреплением, которые можно было бы ускорить с помощью этого подхода.

"Поскольку рассуждения становятся основной нагрузкой, стимулирующей спрос на инференс, работа Цинхао Ху под названием TLT является отличной работой для решения проблемы вычислительного узкого места при обучении этих рассуждающих моделей. Я думаю, что этот метод будет очень полезен в контексте эффективных вычислений ИИ", - говорит Хан.

Эта работа финансируется MIT-IBM Watson AI Lab, программой MIT AI Hardware, MIT Amazon Science Hub, компанией Hyundai Motor Company и Национальным научным фондом.