Исследователи удвоили скорость обучения ИИ, используя простаивающее время GPU

Обучение больших языковых моделей (LLM) — чрезвычайно дорогостоящий процесс. Дело не только в количестве графических процессоров (GPU), но и в эффективности их использования. По мере роста моделей даже небольшие неэффективности оборачиваются огромными затратами времени и энергии.

Теперь команда исследователей из Массачусетского технологического института (MIT) совместно с коллегами, включая специалистов из Nvidia, заявляет, что нашла удивительно практичный способ использовать простаивающие вычислительные ресурсы во время обучения. В некоторых случаях это позволяет сократить общее время обучения почти вдвое.

Проблема, на которую они нацелились, кроется в обучении с подкреплением (RL), особенно на этапе, известном как «прокатка» (rollout). На этом шаге модель генерирует несколько вариантов ответов, чтобы понять, какие действия приводят к лучшим результатам. Это важно для LLM, ориентированных на рассуждения, но процесс очень медленный.

Фактически, этап прокатки может занимать до 85% общего времени выполнения. Виновником является так называемое «длиннохвостое распределение» длин ответов. Большинство сгенерированных ответов завершаются быстро, но небольшое их количество выполняется гораздо дольше среднего. Поскольку GPU должны синхронизироваться, более быстрые процессоры часто простаивают в ожидании отстающих.

Решение команды MIT, названное «Укрощение длинного хвоста» (Taming the Long Tail, TLT), борется с этими потерями напрямую. Вместо того чтобы позволять GPU простаивать во время генерации длинных ответов, TLT использует это время для обучения на лету облегчённой «черновой» модели. Эта меньшая модель непрерывно учится у основной по мере прогресса обучения.

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

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

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

Есть и интересный побочный эффект: непрерывно обучаемая черновая модель сама по себе становится полезным артефактом. Поскольку она обучается параллельно с основной, её можно использовать в качестве эффективной модели для вывода в определённых контекстах.

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

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

Исследование доступно на arXiv.org.