DynaTrain: переключение параллелизма при обучении LLM за секунды

Группа исследователей представила DynaTrain — распределённую систему обучения больших языковых моделей (LLM), способную переключать конфигурацию параллелизма за доли секунды. Работа опубликована на arXiv и описывает подход, решающий проблему динамической адаптации обучения under изменяющиеся ресурсы.

Современное обучение LLM сталкивается с постоянными изменениями: колебания доступных вычислительных ресурсов, сдвиги фаз при RLHF, эластичность кластеров. Традиционные фреймворки статичны и требуют длительного перезапуска или сохранения чекпоинтов. DynaTrain устраняет это узкое место.

В основе DynaTrain лежит абстракция Virtual Parameter Space (VPS) — единое логическое координатное пространство для всех распределённых состояний. Любая конфигурация параллелизма сводится к детерминированному отображению, а сложные переходы — к простым геометрическим пересечениям. Это позволяет унифицировать работу с произвольным многомерным параллелизмом.

Поверх VPS построен уровень маршрутизации и передачи состояний, выполняющий обмен данными между рангами по deadlock-свободному расписанию с учётом памяти. Elastic Device Manager запускает построение нового мира параллельно с продолжающимся обучением, маскируя затраты на смену топологии.

Тесты проводились на плотных и MoE-моделях размером до 235 миллиардов параметров. DynaTrain переконфигурировал 70B плотную модель менее чем за 2 секунды, а 235B MoE-модель — за 4,36 секунды. Это до трёх порядков быстрее, чем существующие системы на основе чекпоинтов и эластичные решения, при сохранении корректности обучения.

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

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