Масштабирование тренировки LLM: четыре вида параллелизма
Почему одну модель тренируют тысячи ускорителей
Современные языковые модели не помещаются в память одного ускорителя: для тренировки трёхмиллиардной модели нужно не меньше 30–60 ГБ видеопамяти, а в процессе обучения требования к памяти вырастают в 10–20 раз. Поэтому обучение распределяют между десятками и сотнями чипов. Главная задача такого масштабирования — чтобы производительность росла пропорционально числу ускорителей, а не упиралась в накладные расходы на синхронизацию. В свежей главе о масштабировании LLM разобраны четыре стратегии параллелизма.
Параллелизм данных — простой старт
Идея предельно проста: батч режется на части и раздаётся ускорителям, каждый хранит полную копию модели и считает свой кусок градиентов. В конце операция AllReduce собирает локальные градиенты в глобальные и обновляет веса на всех устройствах. Применять его стоит всегда, когда модель целиком влезает в ускоритель: масштабируется почти линейно, а синхронизацию можно делать асинхронно. Ограничение одно — память: тренируемый трансформер занимает на порядок больше места, чем финальные веса.
FSDP и ZeRO: без дублирования
Fully-Sharded Data Parallelism, известный как ZeRO-шардинг, устраняет главный недостаток обычного параллелизма данных — избыточность. Вместо полных копий модели веса, градиенты и состояния оптимизатора распределяются по ускорителям: в ZeRO-1 шардируются только состояния оптимизатора, в ZeRO-2 добавляются градиенты, а в ZeRO-3 — и сами веса. Перед умножением матриц нужные веса собираются через AllGather, после обновления — снова раскидываются. Каждое устройство хранит только то, что нужно его шарду — поэтому подход считается основным для моделей среднего размера.
Тензорный параллелизм и конвейеры
Тензорный параллелизм (Megatron-шардинг) меняет оси разбиения матриц и заставляет ускорители обмениваться не весами, а активациями. Эффективен он, когда активаций меньше, чем весов, — чего добиваются комбинацией с FSDP: связка FSDP+TP выходит вычислительно-ограниченной уже при батче примерно в сто токенов. Конвейерный параллелизм, напротив, режет модель по слоям и размещает их на разных GPU. Он пересылает минимум данных и хорошо работает на медленных шинах, но страдает от «конвейерного пузыря» — простоя устройств в ожидании конца цепочки. Лечат его микробатчингом, как в DeepSeek V3.
Выбор стратегии — компромисс между памятью, скоростью шины и простотой реализации. Чаще используют комбинацию методов — так тренировка гигантской модели становится предсказуемым процессом.
А при параллелизме данных AllReduce не превращается в узкое место, когда ускорителей уже несколько сотен? Или на таких масштабах всё равно приходится комбинировать его с другими стратегиями синхронизации?
Хорошо, что такие главы выходят открыто — про четыре вида параллелизма обычно приходится собирать информацию по крупицам из блогов и докладов. У нас связка DeepSpeed и FSDP закрывает почти все задачи до десятков миллиардов параметров, а вот с tensor parallel на нескольких нодах приходится повозиться с коммуникацией. Интересно, в главе есть конкретные цифры по накладным расходам AllReduce на больших кластерах или только общие рекомендации? И ещё вопрос: материалы можно использовать в коммерческих проектах или они только для чтения?
На своей практике добавил бы нюанс: пока модель влезает в одну ноду, связка FSDP с параллелизмом данных масштабируется почти линейно, а вот pipeline parallel стоит включать осторожно — на восьми картах у нас пустые окна в конвейере съедали часть выигрыша по памяти. В итоге остановились на комбинации data и tensor parallel, шаг стал заметно стабильнее, чем с полным 3D-подходом.
@Павел, а пробовали закрывать пустые окна конвейера уменьшенным размером батча на ранних стадиях? У нас на DeepSpeed такой приём почти полностью убрал просадку по памяти, хотя настраивать пришлось дольше, чем хотелось бы.