Оптимизация работы GPU позволила вдвое ускорить обучение ИИ

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

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

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

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

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

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

Источник: iXBT