Как обучали BLOOM: технологии Megatron и DeepSpeed для 176-миллиардной модели

Обучение больших языковых моделей — одна из самых ресурсоёмких задач в современном машинном обучении. В 2022 году команда Hugging Face совместно с исследователями по всему миру представила BLOOM — открытую мультиязычную модель с 176 миллиардами параметров, обученную на 46 языках. Чтобы справиться с колоссальными вычислительными требованиями, разработчики применили комбинацию двух мощных фреймворков: Megatron от NVIDIA и DeepSpeed от Microsoft. Эта статья подробно разбирает, как именно проходило обучение BLOOM, какие технологии оптимизации использовались и почему этот опыт важен для всего сообщества AI.

Что такое BLOOM и почему его обучение стало событием

BLOOM (BigScience Large Open-science Open-access Multilingual Language Model) — это не просто очередная языковая модель. Это результат масштабного международного проекта BigScience, объединившего сотни исследователей. В отличие от многих коммерческих моделей, BLOOM полностью открыт: веса, код обучения и технические отчёты доступны публично. Модель обучена на 46 языках и 13 программных языках, что делает её уникальным инструментом для мультиязычных задач.

Обучение модели такого масштаба потребовало 1.5 миллиона GPU-часов и заняло 3.5 месяца. Для сравнения: GPT-3 (175 миллиардов параметров) обучался на примерно тех же вычислительных мощностях, но детали его обучения остаются закрытыми. Публикация технического обзора от Hugging Face позволяет сообществу заглянуть «под капот» и понять, какие методы оптимизации позволили эффективно распределить нагрузку на 384 GPU NVIDIA A100.

Как работают Megatron и DeepSpeed вместе

Тензорный параллелизм от Megatron

Megatron-LM — это фреймворк от NVIDIA, предназначенный для обучения гигантских моделей. Его ключевая идея — тензорный параллелизм, при котором один слой нейросети разрезается на части и распределяется между несколькими GPU. Например, если у вас есть слой с весами размером 1024x1024, его можно разделить на четыре части по 256x256 и разместить на четырёх GPU. Каждый GPU вычисляет свою часть, а затем результаты объединяются.

При обучении BLOOM тензорный параллелизм использовался для распределения вычислений внутри одного узла сервера. Это особенно эффективно для моделей с огромными скрытыми размерностями (у BLOOM — 14336). Без такого подхода память одного GPU быстро переполнилась бы.

Пайплайновый параллелизм от DeepSpeed

DeepSpeed — библиотека от Microsoft, которая дополняет Megatron. Её главная фишка — пайплайновый параллелизм. Вместо того чтобы делить слои внутри одного GPU, модель разбивается на последовательные стадии, и каждая стадия выполняется на отдельном GPU. Представьте конвейер на заводе: один GPU обрабатывает первые слои, передаёт результат следующему, и так далее.

В BLOOM пайплайновый параллелизм позволил объединить несколько узлов в единую вычислительную цепочку. Это снижает объём межсоединений, так как данные передаются только между соседними стадиями. В комбинации с тензорным параллелизмом это дало синергетический эффект: внутри узла — тензорный параллелизм, между узлами — пайплайновый.

Zero-оптимизация для экономии памяти

Третьим ключевым компонентом стала оптимизация ZeRO (Zero Redundancy Optimizer) от DeepSpeed. ZeRO устраняет избыточное дублирование данных, которое возникает при распределённом обучении. Обычно каждый GPU хранит полную копию параметров модели, градиентов и состояний оптимизатора. Для 176-миллиардной модели это потребовало бы терабайты памяти на каждом устройстве.

ZeRO разделяет эти данные между GPU: каждый хранит только свою часть. Например, ZeRO-3 распределяет параметры, градиенты и состояния оптимизатора по всем GPU. При необходимости данные собираются на лету. Это позволяет обучать модели, которые в десятки раз больше, чем позволяла бы память одного GPU.

Какие ещё оптимизации применялись

Помимо трёх основных техник, команда BLOOM использовала несколько дополнительных приёмов. Во-первых, смешанная точность (mixed precision) — вычисления велись в формате FP16, а некоторые операции — в FP32 для стабильности. Это ускорило обучение и снизило нагрузку на память.

Во-вторых, применялась активационная чекпоинтинг (activation checkpointing): промежуточные значения не сохранялись на каждом шаге, а пересчитывались при обратном распространении. Это позволило сэкономить память ценой небольшого замедления.

В-третьих, использовалась оптимизация коммуникаций: библиотека NCCL от NVIDIA была настроена для максимальной пропускной способности между GPU. Поскольку модель обучалась на кластере из 48 узлов по 8 GPU A100 каждый, скорость обмена данными критически влияла на общее время.

Какие проблемы возникали при обучении

Обучение такой большой модели не обошлось без трудностей. Одна из главных — стабильность. При 176 миллиардах параметров даже малейшая ошибка в градиентах могла привести к расходимости. Команде пришлось тщательно подбирать скорость обучения (learning rate) и использовать градиентный клиппинг.

Другая проблема — аппаратные сбои. За 3.5 месяца непрерывной работы на 384 GPU неизбежно возникали отказы. DeepSpeed предоставил механизмы сохранения состояния (checkpointing), позволяющие возобновлять обучение с последней сохранённой точки. Без этого потеря прогресса была бы катастрофической.

Также пришлось решать задачу балансировки нагрузки. Из-за разной длины последовательностей в мультиязычных данных некоторые GPU могли простаивать. Для этого использовалась динамическая настройка размера батча.

Почему выбор Megatron и DeepSpeed важен для сообщества

Комбинация Megatron и DeepSpeed — не единственный способ обучать большие модели. Например, Google использует собственный фреймворк Pathways, а Meta — Fairscale. Однако открытость BLOOM и детальное описание конфигурации позволяют любому исследователю воспроизвести результаты или адаптировать подход для своих задач.

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

Какие модели можно обучить с помощью этих технологий

Методы, применённые в BLOOM, подходят не только для гигантских моделей. Если у вас есть доступ к нескольким GPU, вы можете использовать тензорный и пайплайновый параллелизм для ускорения обучения моделей среднего размера — например, с 1–10 миллиардами параметров. DeepSpeed ZeRO особенно полезен для моделей, которые не помещаются в память одного GPU, даже если у вас всего 2–4 устройства.

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

Какие технологии используются для обучения больших языковых моделей сегодня

С момента обучения BLOOM появились новые подходы. Например, Microsoft представила DeepSpeed 4, который ещё эффективнее распределяет память. NVIDIA выпустила Megatron-Turing NLG, объединив свои наработки с моделями от Microsoft. Также набирает популярность техника «обучение с нулевым градиентом» (Zero-Offload), когда часть данных выгружается на CPU.

Тем не менее, комбинация Megatron и DeepSpeed остаётся золотым стандартом для открытых проектов. BLOOM доказал, что даже без доступа к суперкомпьютерам можно обучить модель мирового класса, если правильно комбинировать существующие инструменты.

Заключение

Обучение BLOOM стало важной вехой в истории открытого AI. Благодаря публикации технических деталей сообщество получило готовый рецепт для обучения больших языковых моделей. Комбинация Megatron, DeepSpeed и ZeRO позволила эффективно использовать 384 GPU и завершить обучение за 3.5 месяца. Теперь любой исследователь или компания может повторить этот успех, адаптировав методы под свои задачи. Открытость BLOOM — это не просто публикация весов, а передача знаний, которые ускорят развитие всей области.