Глава 4.4 — Распределённое обучение

Содержание

1. От «на чём обучать» к «как физически это обучить»
2. Почему одного ускорителя никогда не хватает на масштабе переднего края
3. Параллелизм данных и его предел
4. Тензорный параллелизм: шардирование самих слоёв
5. Конвейерный параллелизм: последовательное шардирование слоёв
6. ZeRO и FSDP: шардирование состояния оптимизатора вместо модели
7. Память активаций: чекпоинтинг, параллелизм последовательности и перекрытие
8. Комбинирование стратегий: 3D- и 4D-параллелизм
9. Взгляд с точки зрения интервью
10. Вопросы для самопроверки
11. Источники

1. От «на чём обучать» к «как физически это обучить»

Глава 4.3 дала вам способ количественно решить, насколько крупную модель обучать и сколько токенов ей скармливать при заданном бюджете вычислений, и завершилась указанием на то, что она ничего не говорит о том, как вообще разместить такую модель на оборудовании и обучить её за конечное время. Этот пробел — не мелкая деталь реализации: при числе параметров и токенов, которые законы масштабирования теперь регулярно рекомендуют, наивный подход «поместить всю модель на одно устройство и запустить градиентный спуск» не просто медленный — он попросту невозможен, и стоит сформулировать причины этого точно, а не отмахнуться общими словами. Эта глава посвящена системной инженерии, которая вообще делает физически реализуемыми обучающие прогоны масштаба переднего края, независимо от каких-либо алгоритмических решений, рассмотренных ранее.

2. Почему одного ускорителя никогда не хватает на масштабе переднего края

Есть две отдельные причины, по которым единственный ускоритель не может обучить модель масштаба переднего края, и их стоит держать раздельно, поскольку они требуют разных решений. Первая — память: обучение требует хранения не только параметров модели, но и их градиентов (того же размера, что и параметры), а для адаптивных оптимизаторов, таких как Adam, — ещё и состояния оптимизатора для каждого параметра (как правило, два дополнительных тензора того же размера, что и параметры, для скользящих первого и второго моментов), плюс активации, сохраняемые для обратного прохода. Ещё до учёта активаций параметры плюс градиенты плюс состояние оптимизатора для модели с десятками или сотнями миллиардов параметров, хранимые в разрядностях, обычно используемых для стабильного обучения, попросту с большим запасом превышают память любого доступного сегодня отдельного ускорителя — это жёсткое ограничение, а не неудобство производительности, и никакое терпение его не исправит. Конкретно, для обучения в смешанной точности с Adam статья ZeRO авторства Rajbhandari с соавторами показывает, что эти «состояния модели» стоят примерно $16\,\Psi$ байт на параметр, где $\Psi$ — число параметров: 2 байта на fp16-параметр плюс 2 байта на fp16-градиент плюс 4+4+4 байта на fp32 master-копию и два момента оптимизатора Adam. Для модели в 70 млрд параметров это свыше 1 ТБ одних только состояний модели, прежде чем сохранена хотя бы одна активация, — именно поэтому ни один отдельный ускоритель даже близко не вмещает модель масштаба переднего края. Вторая причина — пропускная способность: даже в гипотетическом мире, где всё помещалось бы в память, обучение на числах токенов, которые сейчас рекомендуют законы масштабирования, на одном устройстве заняло бы время, измеряемое годами или десятилетиями, что операционно бесполезно независимо от памяти. Распределённое обучение существует, чтобы решить обе проблемы одновременно — распределить использование памяти между многими устройствами и распределить вычисления между многими устройствами, чтобы свести время выполнения к чему-то практически приемлемому.

3. Параллелизм данных и его предел

Простейшая и старейшая стратегия распределённого обучения — параллелизм данных: реплицируйте всю модель на каждом устройстве, разбейте каждый обучающий батч на непересекающиеся шарды, пусть каждое устройство вычислит прямой и обратный проход на своём собственном шарде независимо, а затем усредните полученные градиенты по всем устройствам перед применением шага оптимизатора, так что все реплики остаются синхронизированными и эффективно ведут себя как один большой обучающий прогон с большим батчем, распределённый по многим устройствам. Это привлекательно, поскольку не требует практически никаких изменений самой модели — каждое устройство хранит идентичную полную копию модели, — и единственная новая инженерная нагрузка — это шаг коммуникации для усреднения градиентов (all-reduce между устройствами) после каждого обратного прохода.

Его предел — в точности та проблема памяти, описанная выше: чистый параллелизм данных требует, чтобы каждое отдельное устройство хранило полную модель, её полные градиенты и её полное состояние оптимизатора, а значит, параллелизм данных сам по себе никак не помогает, если модель изначально не помещается на одном устройстве. Он масштабирует пропускную способность — больше устройств обрабатывают больше частей батча параллельно, — но не масштабирует объём модели, которую вы можете обучить, а значит, на масштабе переднего края его приходится комбинировать со стратегией, которая реально шардирует саму модель по устройствам, а не только данные.

4. Тензорный параллелизм: шардирование самих слоёв

Тензорный параллелизм, иногда называемый параллелизмом модели, решает проблему памяти напрямую, разбивая отдельные весовые матрицы — и построенные из них слои — между устройствами, так что ни одному отдельному устройству никогда не требуется хранить полную копию какого-либо конкретного слоя. Megatron-LM, подход Шойби (Shoeybi) с соавторами к этой проблеме, детально прорабатывает, как разбить два доминирующих по вычислительной нагрузке компонента слоя трансформера — блок внимания и полносвязный (feedforward) блок — между устройствами таким образом, что требуется лишь небольшой, чётко определённый объём коммуникации на слой, а не коммуникация после каждого отдельного матричного умножения. Для полносвязного блока это обычно означает разбиение весовой матрицы первого линейного слоя по столбцам между устройствами (так что каждое устройство независимо вычисляет свой срез промежуточной активации, не требуя коммуникации на этом шаге) и разбиение весовой матрицы второго линейного слоя по строкам (так что суммирование частичных выходов со всех устройств — единственный шаг коммуникации — восстанавливает правильный итоговый выход). Внимание разбивается аналогичным образом, путём распределения самих голов внимания между устройствами, поскольку вычисление каждой головы независимо от других вплоть до финальной проекции выхода.

Центральный компромисс здесь — стоимость коммуникации: тензорный параллелизм требует синхронизации устройств (как правило, через all-reduce) в определённых точках внутри каждого отдельного слоя, а не только один раз за батч, как это делает параллелизм данных, а значит требует очень высокопропускных, низколатентных межсоединений между задействованными устройствами — он хорошо работает внутри одного высокопропускного сервера или небольшого кластера таких серверов, но быстро деградирует, если попытаться растянуть его на устройства, соединённые более медленными сетями, потому что частая коммуникация на каждом слое становится узким местом, а не сами вычисления. Именно поэтому тензорный параллелизм обычно применяется с относительно небольшой «шириной» — разбиением на несколько устройств или несколько десятков устройств внутри плотно связанного сетью узла или стойки, — а не как единственная стратегия распределения модели по целому крупному кластеру.

5. Конвейерный параллелизм: последовательное шардирование слоёв

Конвейерный параллелизм подходит к той же проблеме памяти иначе: вместо разбиения отдельных слоёв между устройствами он разбивает слои модели на непрерывные группы (стадии) и назначает каждую стадию отдельному устройству, так что данные проходят через устройства последовательно, как на сборочной линии — устройство один вычисляет первые несколько слоёв и передаёт свой выход устройству два, которое вычисляет следующие несколько слоёв, и так далее. Это требует гораздо меньше коммуникации на шаг, чем тензорный параллелизм, поскольку передавать нужно только активации на границах стадий, но вносит другую проблему: наивно только одно устройство выполняет полезную работу в любой момент времени, пока остальные простаивают в ожидании прихода своего входа, что сводит на нет весь смысл параллелизации.

Стандартное решение — микробатчинг: разбить каждый обучающий батч на несколько меньших микробатчей и подавать их в конвейер один за другим по смещённому расписанию, так что пока устройство два обрабатывает более позднюю стадию первого микробатча, устройство один уже обрабатывает более раннюю стадию второго микробатча, поддерживая занятость всех устройств одновременно разными микробатчами на разных этапах конвейера. Неизбежно остаётся период запуска и опустошения конвейера — «пузырь» (bubble), — во время которого конвейер заполняется или опустошается и не все устройства полностью загружены, и значительная часть инженерии конвейерного параллелизма заключается в выборе числа и размера микробатчей, а также конкретного расписания прямых и обратных проходов между ними, чтобы минимизировать долю общего времени обучения, теряемую на этот пузырь.

6. ZeRO и FSDP: шардирование состояния оптимизатора вместо модели

ZeRO, Zero Redundancy Optimizer, описанный Раджбхандари (Rajbhandari) с соавторами, использует подход, действительно ортогональный предыдущим: вместо того чтобы вообще разбивать слои модели или весовые матрицы между устройствами, он сохраняет простую структуру параллелизма данных — каждое устройство обрабатывает свой шард данных, — но устраняет избыточность хранения каждым устройством идентичной полной копии состояний оптимизатора, градиентов и (на своей самой агрессивной стадии) самих параметров. Обычный параллелизм данных дублирует всё это на каждой реплике, хотя в любой конкретный момент каждому устройству строго необходима лишь та часть, которая относится к конкретному вычислению, выполняемому им прямо сейчас; ZeRO вместо этого шардирует эти тензоры между группой устройств параллелизма данных, так что каждое устройство хранит лишь долю состояния оптимизатора, градиентов или параметров, динамически собирая нужные ему части от других устройств через коммуникацию именно тогда, когда конкретное вычисление требует полного тензора, и снова освобождая их после этого. Это резко сокращает память на устройство по сравнению с обычным параллелизмом данных — часто почти в число раз, равное числу устройств параллелизма данных, в зависимости от используемой стадии ZeRO: на самой агрессивной стадии ZeRO память на устройство под состояния модели падает примерно до $\dfrac{16\,\Psi}{N_{dp}}$ байт, где $N_{dp}$ — число устройств параллелизма данных, вместо полных $16\,\Psi$, которые платит каждое устройство при обычном параллелизме данных, — не требуя архитектурного разбиения самой модели, как этого требуют тензорный или конвейерный параллелизм, что делает его сравнительно простым в применении к уже существующей реализации модели.

FSDP, Fully Sharded Data Parallel, — это реализация по существу той же идеи, которая стала стандартной, широко используемой на практике реализацией шардирования в стиле ZeRO, встроенной в основные фреймворки обучения, а не существующей только в отдельной специализированной библиотеке. Концептуальное содержание то же, что и у наиболее агрессивной стадии ZeRO — шардировать параметры, градиенты и состояние оптимизатора между воркерами параллелизма данных, собирая полные шарды параметров непосредственно перед вычислением для каждого слоя и немедленно освобождая их после, — но стоит знать FSDP именно по имени, поскольку это термин, наиболее часто используемый в современной документации фреймворков обучения и в разговорах об этой технике на интервью.

7. Память активаций: чекпоинтинг, параллелизм последовательности и перекрытие

Всё в разделах 3–6 сокращает память, занимаемую параметрами, градиентами и состоянием оптимизатора. Есть четвёртый потребитель памяти устройства, которого ни одна из них не касается: активации — промежуточные выходы каждого слоя, вычисляемые на прямом проходе и сохраняемые, потому что они нужны обратному проходу для вычисления градиентов. При размерах батча и длинах последовательности, реально используемых в обучении переднего края, память активаций — не погрешность округления: для достаточно глубокой модели на длинных последовательностях она может превысить память, занимаемую самими параметрами, и именно поэтому конфигурация, комфортно помещающая свои веса и состояние оптимизатора под ZeRO/FSDP, всё равно может упереться в нехватку памяти, стоит вам увеличить длину последовательности или размер батча.

Чекпоинтинг активаций (также называемый gradient checkpointing, по формулировке Chen с соавторами) напрямую обменивает эту память на вычисления: вместо того чтобы сохранять активации каждого слоя на протяжении всего прямого прохода, сохраняется лишь разреженный набор «контрольных точек» на выбранных слоях, остальное отбрасывается, а на обратном проходе отброшенные активации пересчитываются по требованию, заново прогоняя прямое вычисление от ближайшей сохранённой контрольной точки. Chen с соавторами показали, что это позволяет памяти масштабироваться примерно как $O(\sqrt{L})$ вместо $O(L)$, где $L$ — число слоёв, ценой одного дополнительного прямого прохода через пересчитываемые сегменты — обычно оцениваемого примерно в 30% дополнительных вычислений за существенное снижение пикового потребления памяти активациями. Это учебный пример компромисса, к которому эта книга постоянно возвращается: потратить больше FLOPs, которых сравнительно много, чтобы сэкономить память, которая на масштабе переднего края является более узким ограничением.

Параллелизм последовательности, описанный Korthikanti с соавторами, атакует конкретный пробел, который тензорный параллелизм (раздел 4) оставляет открытым: такие операции, как LayerNorm и dropout, не разбиваются тензорным параллелизмом в стиле Megatron, поскольку они работают независимо для каждого токена, а не смешивают информацию по скрытому измерению так, как это делают внимание и полносвязный блок, — так что каждое устройство тензорного параллелизма в итоге избыточно хранит полную активацию для этих слоёв, независимо от того, на сколько частей разбиты сами матричные умножения. Параллелизм последовательности вместо этого шардирует именно эти активации по измерению последовательности, внутри той же группы тензорного параллелизма, которая уже обменивается данными ради разбитых матричных умножений, так что избыточные копии исчезают, не добавляя никакого нового паттерна коммуникации сверх того, за что тензорный параллелизм уже платит. Korthikanti с соавторами показывают, что сочетание параллелизма последовательности с более избирательной формой чекпоинтинга активаций — пересчётом только тех конкретных операций, которые дёшево повторить, а не чекпоинтингом целых блоков трансформера — сокращает память активаций существенно сильнее, чем любая из техник по отдельности, добавляя заметно меньше накладных расходов на пересчёт, чем чекпоинтинг каждого блока.

Последний, чисто системный рычаг стоит назвать отдельно, поскольку он ничего не стоит алгоритмически: перекрытие коммуникации с вычислением. Каждая стратегия этой главы вносит какую-то коммуникацию — all-reduce градиентов в параллелизме данных, обмен активациями и градиентами на границах тензорного и конвейерного параллелизма, сбор параметров в FSDP, — и наивная реализация просто ждёт завершения каждого шага коммуникации, прежде чем продолжить, оставляя ускоритель простаивающим, пока данные передаются по сети. Поскольку коммуникация для одной части модели (скажем, слоя, чей обратный проход только что завершился) не зависит от вычисления для другой части (обратного прохода следующего слоя, уже готового начаться), фреймворк может выдать коммуникацию асинхронно и дать ей выполняться в фоне, пока ускоритель продолжает вычислять, эффективно скрывая задержку коммуникации за полезной работой всякий раз, когда между ними нет настоящей зависимости по данным. Современные фреймворки обучения делают это по умолчанию везде, где это позволяет структура зависимостей, и это одна из причин, по которой грамотно спроектированный 3D/4D-параллелизм достигает пропускной способности, гораздо более близкой к теоретическому пределу, ограниченному вычислениями, чем наивная, неперекрывающаяся реализация тех же стратегий.

8. Комбинирование стратегий: 3D- и 4D-параллелизм

На практике ни одна из этих стратегий никогда не используется в одиночку на масштабе переднего края — каждая решает свою собственную проблему узкого места, и реальные крупные обучающие прогоны комбинируют несколько стратегий одновременно, что обычно называют 3D-параллелизмом (параллелизм данных, тензорный и конвейерный, скомпонованные вместе) или 4D-параллелизмом, когда шардирование в стиле ZeRO добавляется как четвёртая, ортогональная ось поверх трёх остальных. Типичная крупномасштабная схема может использовать тензорный параллелизм внутри плотно связанного сетью узла для разбиения отдельных слоёв между несколькими устройствами, конвейерный параллелизм между группами узлов для разбиения глубины модели на стадии, параллелизм данных между многими такими репликами конвейера для одновременной обработки большей части батча, а также шардирование в стиле ZeRO/FSDP внутри измерения параллелизма данных для дальнейшего сокращения избыточного использования памяти. Общее число устройств, которое потребляет такая схема, — это просто произведение размеров, выбранных для каждой оси: $N_{\text{total}} = N_{tp} \times N_{pp} \times N_{dp}$ (с дополнительным множителем на степень шардирования ZeRO/FSDP при построении 4D-параллелизма). Правильно подобрать эту комбинацию — решить, сколько устройств выделить на каждую ось и в какой топологии относительно реальной физической сети, — само по себе является существенной и значимой инженерной дисциплиной, и именно это конкретное, лишённое всякого гламура обстоятельство делает обучение модели переднего края в такой же мере задачей распределённых систем, в какой и задачей машинного обучения.

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

9. Взгляд с точки зрения интервью

В: Почему нельзя просто обучить модель с 70 миллиардами параметров на одном ускорителе, даже располагая неограниченным временем? Сильный ответ разделяет ограничение по памяти и ограничение по пропускной способности: параметры, градиенты и (для оптимизаторов Adam-типа) состояние оптимизатора на параметр вместе превышают память любого отдельного устройства при таком масштабе независимо от того, сколько вы готовы ждать, так что это жёсткий предел по объёму, а не просто проблема скорости — хотя проблема скорости (годы времени выполнения на одном устройстве) тоже реальна и потребовала бы отдельного решения, даже если бы память не была проблемой.

В: В чём реальная разница между тензорным и конвейерным параллелизмом и когда стоит выбирать один из них? Сильный ответ формулирует, что тензорный параллелизм разбивает отдельные весовые матрицы/слои между устройствами так, что каждое устройство хранит шард каждого слоя, требуя частой, высокопропускной коммуникации внутри вычисления каждого слоя, — подходит для плотно связанных сетью устройств, обычно внутри одного узла. Конвейерный параллелизм разбивает слои модели на последовательные стадии между устройствами, требуя коммуникации только на границах стадий, но рискуя простоем («пузырём»), смягчаемым микробатчингом, — лучше подходит для менее тесно связанных групп устройств, например между узлами.

В: Объясните, что ZeRO/FSDP делают иначе, чем тензорный или конвейерный параллелизм. Сильный ответ уточняет, что ZeRO/FSDP вообще не разбивают модель архитектурно — они сохраняют структуру параллелизма данных (каждое устройство обрабатывает разные данные), но устраняют избыточные полные копии состояния оптимизатора, градиентов и параметров, которые обычный параллелизм данных иначе хранил бы на каждом устройстве, шардируя эти тензоры между группой параллелизма данных и собирая их непосредственно перед использованием. Это ортогональная ось экономии памяти, а не замена тензорному/конвейерному параллелизму.

В: Что такое «пузырь» в конвейерном параллелизме и как с ним борются? Сильный ответ описывает пузырь как время простоя в начале и конце обработки батча, когда не все стадии конвейера имеют работу, потому что конвейеру нужно заполниться и опустошиться, и объясняет, что микробатчинг — разбиение батча на более мелкие фрагменты, подаваемые в конвейер по смещённому, перекрывающемуся расписанию — поддерживает занятость большего числа устройств одновременно и снижает (не устраняя полностью) долю времени, теряемого на пузырь.

В: Почему реальные обучающие прогоны переднего края комбинируют несколько стратегий параллелизма, а не выбирают только одну? Сильный ответ объясняет, что каждая стратегия решает своё собственное ограничение ресурса (дублирование памяти, размер слоя, глубина модели, пропускная способность коммуникации), и ни одна из них в одиночку не достаточна на масштабе переднего края — один только тензорный параллелизм не масштабируется через низкопропускные соединения, один только конвейерный параллелизм оставляет избыточность памяти и неэффективность пузыря, один только параллелизм данных не снижает память модели на устройство, — так что их комбинирование (3D/4D-параллелизм) позволяет подобрать размер каждой оси под то ограничение, для которого она лучше всего подходит, например тензорный параллелизм внутри быстро связанного сетью узла и конвейерный/параллелизм данных между узлами.

В: Вы зашардировали параметры и состояние оптимизатора через FSDP, но всё равно упираетесь в память при увеличении длины последовательности. В чём дело и что бы вы сделали? Сильный ответ определяет, что шардирование в стиле FSDP решает проблему параметров, градиентов и состояния оптимизатора, но не памяти активаций, которая растёт с длиной последовательности и размером батча независимо от них. Решение — чекпоинтинг активаций (пересчитывать отброшенные активации на обратном проходе вместо хранения всех) и, если используется тензорный параллелизм, параллелизм последовательности, чтобы убрать избыточные подевайсные копии активаций LayerNorm/dropout, которые один только тензорный параллелизм оставляет дублированными.

В: Что реально даёт «перекрытие коммуникации с вычислением» и почему это не происходит автоматически? Сильный ответ объясняет, что каждая стратегия параллелизма вносит коммуникацию (all-reduce градиентов, обмен активациями, сбор параметров), а наивная реализация блокируется на каждом шаге коммуникации, оставляя ускоритель простаивающим. Там, где часть коммуникации не имеет настоящей зависимости по данным от вычисления, которое могло бы выполняться параллельно (например, all-reduce градиентов одного слоя против обратного прохода следующего), фреймворк может выдать её асинхронно и скрыть задержку за полезным вычислением — реальный выигрыш в пропускной способности, но только там, где граф зависимостей действительно допускает такую перестановку.

10. Вопросы для самопроверки

1. Назовите две независимые причины, по которым единственный ускоритель не может обучить модель масштаба переднего края, и объясните, почему они требуют разных видов решений.
2. Почему параллелизм данных масштабирует пропускную способность, но не максимальный размер модели, которую вы можете обучить?
3. Проследите, как Megatron-LM разбивает два линейных слоя полносвязного блока между устройствами, и объясните, почему разбиение первого по столбцам, а второго по строкам требует лишь одного шага коммуникации, а не одного после каждого матричного умножения.
4. Какую конкретную проблему вносит конвейерный параллелизм, которую не создаёт тензорный параллелизм, и каково стандартное решение?
5. В каком точном смысле ZeRO/FSDP «ортогональны» тензорному и конвейерному параллелизму, а не являются конкурирующей альтернативой им?
6. Какая из стадий ZeRO, как вы ожидаете, экономит больше всего памяти на устройство, и чего это будет стоить взамен?
7. Набросайте правдоподобную схему 3D- или 4D-параллелизма для обучения очень крупной модели на множестве узлов и обоснуйте, какую стратегию параллелизма вы бы назначили на самое быстрое сетевое измерение вашего кластера и почему.
8. Почему память активаций может превысить память параметров даже после шардирования в стиле ZeRO/FSDP, и что чекпоинтинг активаций обменивает ради её снижения?
9. Объясните конкретно, какую избыточность параллелизм последовательности убирает, оставленную одним тензорным параллелизмом, и почему её устранение не требует нового паттерна коммуникации.

11. Источники

• Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., & Catanzaro, B. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053. https://arxiv.org/abs/1909.08053
• Rajbhandari, S., Rasley, J., Ruwase, O., & He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC20. arXiv:1910.02054. https://arxiv.org/abs/1910.02054
• Rasley, J., Rajbhandari, S., Ruwase, O., & He, Y. (2020). DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters. KDD 2020. Not on arXiv. https://www.deepspeed.ai/ (DOI: 10.1145/3394486.3406703)
• Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training Deep Nets with Sublinear Memory Cost. arXiv:1604.06174. https://arxiv.org/abs/1604.06174
• Korthikanti, V., Casper, J., Lym, S., et al. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198. https://arxiv.org/abs/2205.05198