Раcпределенное обучение LLMs

Глобальный рынок «больших языковых моделей» (LLM) переживает стремительный рост. По прогнозам, к концу 2024 года его объем составит 6,4 миллиардов долларов США, а к 2030 году эта сумма достигнет 36 миллиардов долларов.  Быстро растут и технические параметры самих LLM, для работы которых концентрируются все большие вычислительные мощности. Этот процесс сделал распределенный тип «машинного обучения» (ML) не просто выгодным, но необходимым. О нюансах распределенных архитектур мы поговорили с Ольгой Ковалевой, Старшим Аналитиком — разработчиком в команде «Яндекса». Статья ниже – пересказ полученных от нее разъяснений.

Современные LLM c миллиардами и триллионами параметров не могут управляться одним процессором (GPU). Требуются большие группы устройств и сложные архитектуры распределенного «машинного обучения».

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

Распределенное обучение: проблемы масштабирования

Обучение LLM является нетривиальной задачей с точки зрения компьютерных расчетов. Простой пример. Модель GPT-3 включает 175 миллиардов параметров. Ей требуются сотни гигабайт, чтобы просто хранить параметры модели. Если речь идет обо всем объеме того, что нужно для «машинного обучения» при работе с моделью, то потребности памяти оказываются огромны. «Обсчет» состояний оптимизатора, градиентов, активаций, то потребление памяти вполне может превысить 1 террабайт данных. Это существенно превышает возможности любого персонального GPU (графического процессора). Потребуется архитектура, задействующая одновременно десятки и сотни физических устройств, чтобы LLM могла нормально работать.

Помимо количества памяти поражают временные ресурсы, которые нужно выделить. Например, при некоторых алгоритмах, на которых обучается GPT-3 расчетное время завершения обучения – тысячи лет. Естественно, группам разработчиков приходится искать способы ускорить процесс. Распараллеливание алгоритмов обучения LLM становится приоритетной задачей, если речь заходит о сроках и стоимости разработок.

GPT-3 – выдающийся по своим возможностям и техническим параметрам проект. В среднем, разработчикам приложений на основе LLM и «машинного обучения» в технологических компаниях приходится иметь дело с 10-100 миллиардами параметров. Однако даже при таком, более скромном размахе, не удается обойтись без распределенного обучения, и правильная его организация оказывает огромное влияние на результат. Типичный сервер, даже при скромном, если сравнивать с GPT-3, размахе проекта, включает 8 чипов H100 с 80 ГБ памяти каждый. Как же организовать их взаимодействие эффективно?

Распределенное обучение: стратегии распараллеливания

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

1. Прямой проход по пакетам данных с наблюдениями.
2. Обратный проход для расчета градиентов с использованием активаций, накопленных на первом шаге.
3. Обновление параметров с использованием градиентов, полученных на втором шаге.

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

Рассчитывать градиенты при параллелизме данных можно по разному. Есть две основные стратегии в этом вопросе:

• синхронное распараллеливание данных: необходимо убедиться, что выполняется градиентная агрегация после обработки устройствами каждой порции данных для обучения. С одной стороны, такая стратегия позволяет сделать результат обучения более стабильным, однако обратной стороной становится проблема зависимости машин в системе друг от друга. В результате процесс двигается со скоростью самого медленного сервера в группе;
• асинхронное распараллеливание данных: синхронизация градиента в шаге 2) не выполняется и шаг 3) выполняется с группами градиентов от разных GPU, участвующих в процессе. Асинхронное распараллеливание данных применяется не так часто, поскольку выигрыш по производительности и эффективности в сравнении с синхронным распараллеливанием происходит только тогда, когда процесс синхронизации градиентов в шаге 2) по каким-то причинам оказывается неадекватно «тяжелым» процессом. Обычно это не так.

Особняком стоит такой подход, как «накопление градиента». Он не является самостоятельной альтернативой синхронному или асинхронному подходу, а часто дополняет какой-либо из них. «Накопление градиента» заключается в том, что шаги 1) и 2) повторяются много раз, а шаг 3) имеет дело с результатом многократного выполнения этих, предварительных шагов. Метод «накопления градиента» имеет как свои плюсы, так и минусы. Из минусов: приходится устанавливать локальный размер партии наблюдений равным «1». В противном случае, по ряду технических причин, о которых стоит говорить отдельно, теряется независимость порций данных в накоплении градиентов. Также возрастает риск перерасхода памяти.

Распределенное обучение: распараллеливание модели

Если применяется распараллеливание самой LLM. Копируются уже не данные для обучения, а сама модель. Тем самым она распределяется между устройствами, если является слишком большой для одного устройства. Подход требует скрупулезной координации прямых и обратных проходов.

Применения методики распараллеливания модели стараются избегать, поскольку часто при ее использовании не удается избежать высоких коммуникационных издержек производительности между устройствами. Прежде, чем смотреть в эту сторону, лучше рассмотреть альтернативные возможности. Например, многие проблемы может решить замена оборудования. Эффективность современных, промышленных моделей глубинного обучения зависит от количества параметров, которые используются. Она росла в последние годы за счет развития используемого «железа». Вот как в последнее время менялся типичный узел оборудования для ML-проектов:

• 2022 год: GPU A100 c 40 ГБ памяти;
• 2024 годы: 8 чипов H100 с 80 ГБ памяти каждый;
• 2025-2026 (прогноз): 8 графических процессоров B200 с 160 ГБ памяти каждый.

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

Подход, связанный с параллелизмом модели так же, как и в случае с параллелизмом данных имеет две реализации:

• тензорный параллелизм: делит между отдельными физическими устройствами операции модели, например, умножение матриц. Если LMM выстроена на архитектуре трансформеров, то между различными машинами узла делятся головки внимания или измерений сети;
• кластерное обучение (пайплайн-параллелизм): модель разбивается не по операциям, а по слоям. Активации и частичные производные передаются между GPU.

Оба подхода имеют свои достоинства и недостатки и применяются в зависимости от специфических особенностей проекта. Если при пайплайн-параллелизме деление по слоям плохо скоординировано, то в процессе обучения модели возникают «бутылочные горлышки», на которых теряется эффективность модели. Методики планирования микропакетов и Bubble Minimization позволяют исправить ситуацию. Первые предназначены для того, чтобы исследовать, как данные передаются от машины к машине. Группа методов «Минимизация пузырьков» позволяет решить возникающие проблемы. Например, алгоритм 1F1B очень эффективно компенсировать ситуацию, когда «бутылочное горлышко» возникло из-за задержки в одном слое, обрабатываемом одной машиной. Целый ряд проблем помогают решить методы GPipe и PipeDream.

Распределенное обучение: гибридные подходы

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

• начинаем с параллелизма данных, не используем ничего другого, пока размер локального пакета не достигнет «1» с помощью накопления градиента;
• вводим контрольные точки активации;
• разделяем модель внутри узла и сравниваем тензорный параллелизм с конвейеризацией.

После всех этих шагов в нашем распоряжении 1280 ГБ доступной памяти, что позволяет хранить полтриллиона параметров в fp16. После этого применения таких противоречивых мер, как параллелизм модели может уже не потребоваться.

Распределенное обучение: стратегии по оптимизации утилизации памяти

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

• контрольные точки: выборочное создание контрольных точек сохраняет только определенные активации при прямом проходе и пересчитывает остальные при обратном проходе. Это позволяет экономить вычислительные ресурсы ради памяти. Стратегии градиентных контрольных точек должны обеспечивать баланс между экономией памяти и вычислительными затратами. Распространенные стратегии включают установку контрольных точек через каждые N слоев или использование более сложных алгоритмов, управляемо оптимизирующих соотношение вычислений и памяти.
• шардинг состояний: шардинг состояний оптимизатора распределяет его состояния (например, дисперсию Адама) по устройствам, чтобы снизить требования к памяти на каждое устройство. Этот метод необходим для обучения очень больших моделей. Оптимизатор с нулевой избыточностью (ZeRO) устраняет избыточность памяти, распределяя состояния оптимизатора, градиенты и даже параметры по устройствам. ZeRO-3 может снизить требования к памяти на порядки.

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

Распределенное обучение: координация работы нескольких узлов

Планирование заданий на нескольких узлах технический устройств требует координации с планировщиками кластера, такими как Slurm или Kubernetes. Задания на обучение должны корректно обрабатывать сбои GPU, участвующих в обучении. При распределении ресурсов необходимо учитывать как вычислительные ресурсы, так и пропускную способность сети. Неоптимальное распределение может создавать узкие места в коммуникациях, которые серьезно влияют на производительность.

Механизмы обеспечения стабильной работы необходимы для длительных задач обучения.

Распределенное обучение: архитектура сети

Высокоскоростные коммуникационные технологии, такие как NVLink и RoCE помогают решать многие проблемы распределенного обучения:

• NVLink. Это высокоскоростной интерфейс связи соединяет процессоры внутри узла устройств. Обладает очень высокой пропускной способностью – до 900 ГБ/с. Благодаря этому достигается управляемый параллелизм моделей внутри узла;
• RoCE. Это технология удаленного прямого доступа к памяти. поскольку он работает через стандартный кабель Ethernet.

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

Распределенное обучение: популярные фреймворки

Статья была бы неполной, если бы не поговорили о популярных библиотеках, оркестрирующих распределенное обучение. Вот некоторые из них:

• PyTorch Distributed: предоставляет гибкие примитивы для реализации различных стратегий распараллеливания. Такие функции, как DistributedDataParallel (DDP) и Fully Sharded Data Parallel (FSDP) упрощают распространенные шаблоны обучения LLM;
• DeepSpeed: ​​предлагает сложные методы оптимизации, включая оптимизатор ZeRO, конвейерный параллелизм и 3D-параллелизм. Он особенно хорошо подходит для крупномасштабного обучения LLM;
• FairScale: предоставляет модульные компоненты для распределенного обучения, включая реализации конвейерного параллелизма и методы оптимизации памяти;
• Megatron-LM: демонстрирует лучшие практики обучения очень больших моделей трансформаторов с оптимизированными реализациями тензорного и конвейерного параллелизма.

Наиболее развитым фреймворком является, конечно, PyTorch, однако в зависимости от ситуации другие библиотеки могут оказаться более предпочтительным выбором.

Распределенное обучение: рекомендации по реализации

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

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

Распределенное обучение: оптимизация производительности

Как в других аспектах, здесь есть несколько подходов. Перечислим некоторые из них:

• профилирование связи: выявляет узкие места в синхронизации градиента и передаче активации. Такие инструменты, как Nsight systems очень удобны для поиска узких мест в производительности. Nsight systems – это распространенный профилировщик ядра CUDA;
• отслеживание использования памяти на всех устройствах: например, существует команда nvidia-smi (связанная с одноименной утилитой для командной строки). Использовав ее можно посмотреть самую разную статистику: текущее использование памяти, температуру графического процессора и многое другое. В больших кластерах необходимо собирать эту информацию и отслеживать ее по всему кластеру с помощью аналитических инструментов, таких как Grafana;
• оптимизация пропускной способности: на практике постоянно приходится балансировать между вычислениями на GPU и сетевыми коммуникациями. Правильно настроенная распределенная конфигурация должна купировать большую часть сетевых издержек.

Оптимальная конфигурация зависит от архитектуры модели, размера пакета и характеристик оборудования.

Распределенное обучение: эффективность масштабирования

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

Анализ накладных расходов на коммуникацию помогает определить точку убывающей отдачи при добавлении устройств. Оптимальное количество устройств зависит от размера модели, размера пакета и характеристик сети.

Распределенное обучение: распространенные ошибки

Дисбаланс нагрузки между устройствами или этапами конвейера может серьезно повлиять на производительность. Тщательное профилирование и балансировка крайне важны. Фрагментация памяти может привести к ошибкам нехватки памяти даже при достаточном объеме памяти. В принципе, фрагментация памяти – контролируемый процесс, но он спрятан на низком уровне системы. По этой причине задача обычно автоматизирована с помощью CUDA и фреймворков глубокого обучения. Другое дело, что любая автоматика может дать сбой, поэтому нужен эффективный мониторинг. Следите также за пулом памяти и пересчетом активаций.

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

Распределенное обучение: отладка

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

Современные ML-приложения предотвращают строительные катастрофы за счет создания «цифровых двойников» зданий, проводят хирургические операции и проводят хирургические операции (робот Da Vinci), поэтому права на ошибку у специалистов по ML зачастую нет. Без эффективной отладки просто нельзя обойтись.

Распределенное обучение: будущие направления развития

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

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

Распределенное обучение: развитие оборудования

Специализированное оборудование на основе «искусственного интеллекта», такое как тензорные процессоры и специализированные чипы для обучения требует адаптации стратегий распараллеливания к своим уникальным характеристикам. Его эволюция может создать новые стратегии реализации распределенного обучения LLM или сделать неактуальными старые.

Распределенное обучение: «софт»

Интеграция фреймворков движется в сторону более прозрачной и автоматической обработки сложной распределенной обучающей среды. Облачные решения предоставляют управляемые распределенные обучающие сервисы, которые значительной частью сложной инфраструктуры управляют «под капотом», оставляя разработчикам готовые, более простые и общие методы.

Выводы

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