85% мощности суперкомпьютеров сгорало впустую. Но крошечная нейросеть одолела простои видеокарт и ускорила тренировку ИИ в 3 раза

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

Исследователи из Массачусетского технологического института вместе с коллегами из компании NVIDIA, Швейцарской высшей технической школы Цюриха, лаборатории MIT–IBM Watson AI Lab и Университета Массачусетса в Амхерсте предложили способ использовать этот простой для ускорения обучения. Метод ориентирован на модели рассуждения, которые тренируют через обучение с подкреплением. По словам авторов, система ускоряет процесс без потери точности. Результаты планируют представить на конференции ASPLOS 2026, которая пройдёт в Питтсбурге 22–26 марта.

Узкое место скрыто в самой схеме обучения с подкреплением. Чтобы развить у модели привычку проверять свои шаги, ей дают запрос и просят сгенерировать несколько вариантов ответа. Затем система выбирает лучший вариант, начисляет ему награду и обновляет параметры модели с учётом результата. Такой цикл повторяется тысячи раз, постепенно смещая модель к более надёжным стратегиям рассуждения.

Авторы отмечают, что основное время уходит не на обновление параметров, а на получение множества вариантов ответа. В статье эту стадию называют rollout, по сути это многократный прогон генерации, когда модель разворачивает ответы на один и тот же запрос, чтобы затем было из чего выбирать. По оценке исследователей, на этот этап может приходиться до 85% времени выполнения при обучении с подкреплением, тогда как пересчёт весов занимает сравнительно небольшую долю.

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

Команда взяла приём, который применяют для ускорения генерации, и адаптировала его под обучение с подкреплением. Приём называется speculative decoding (спекулятивное декодирование). Схема строится вокруг пары моделей: рядом с большой работает меньшая и более быстрая. Малая модель быстро предлагает продолжения ответа, а большая проверяет эти предложения и принимает те, которые совпадают с тем, что она сама сгенерировала бы при обычной генерации. Проверка идёт сразу пачкой, поэтому большая модель тратит меньше времени на последовательное дописывание текста шаг за шагом.

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

Чтобы обойти это ограничение, исследователи разработали систему Taming the Long Tail, сокращённо TLT. Название отсылает к проблеме длинного хвоста, когда редкие длинные вычисления тормозят весь процесс. Система состоит из двух взаимосвязанных частей, обе работают автоматически и подстраиваются под нагрузку.

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

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

Модель-помощник сделали лёгкой специально, чтобы её можно было быстро дообучать. Кроме того, система повторно использует часть того, что уже делается в процессе обучения большой модели, и применяет это же для подстройки малой. За счёт такого переиспользования получается дополнительный выигрыш по времени.

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

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

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