[Перевод статьи] ZeroRO сегментирования распределенных параметров обучения

0x00 сводка

Разделение параметров заключается в разделении параметров модели на каждый графический процессор, чтобы достичь цели использования меньшего количества графических процессоров для проведения крупномасштабного обучения модели. В этой серии будет проанализировано разделение параметров с помощью документов, блогов и кода Google, Microsoft и Facebook. Всего будет от 5 до 6 статей.

Эта статья начинается сZeRO: Memory Optimizations Toward Training Trillion Parameter ModelsиDeepSpeed: Extreme-scale model training for everyoneЭтот оптимизатор, разработанный Microsoft в основном для анализа, может эффективно использовать видеопамять, которая распределяет величины состояния модели (состояние оптимизатора, градиенты и параметры модели) на несколько параллельных графических процессоров.Обучайте модели с многомиллиардными параметрами с помощью параллелизма моделей.

ZeRO представляет собой комбинацию двух методов: ZeRO-DP и ZeRO-R. ZeRO-DP — это усовершенствованный механизм параллелизма данных, который использует стратегию динамической связи для разделения состояний, градиентов и параметров оптимизатора, чтобы свести к минимуму связь и избежать избыточности состояний модели. ZeRO-R использует активируемый разделами пересчет, буферы постоянного размера и механизмы динамической дефрагментации памяти для оптимизации потребления памяти для оставшихся состояний.

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

Другие статьи из этой серии:

[Анализ исходного кода] Распределенный PyTorch ZeroRedundancyOptimizer

0x01 Обзор

Эта часть в основном переведена сDeepSpeed: Extreme-scale model training for everyone.

1.1 Проблемы

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

1.1.1 Эффективность видеопамяти

Память, необходимая для обучения модели с триллионом параметров, намного превышает объем памяти одного графического процессора. Например, для обучения смешанной точности с оптимизатором Adam требуется около 16 ТБ видеопамяти для хранения состояния модели (параметров, градиентов и состояния оптимизатора). Только для хранения состояния модели требуется 400 графических процессоров Nvidia A100 (по 40 ГБ видеопамяти каждый).

Функция активации также должна занимать дополнительную видеопамять, которая увеличивается с размером пакета. Если размер пакета равен 1, для обучения модели с триллионом параметров потребуется более 1 ТБ видеопамяти для хранения активаций. Люди также пытаются использоватьcheckpointЧтобы активировать процесс памяти, используется для расчета изменения памяти, память которой может быть уменьшена примерно до 20 ГБ, но для целей обучения потребность в памяти все еще слишком высока.

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

1.1.2 Вычислительная эффективность

Исследования на основе OpenAIlaw of scalingДля оценки требуется около 5000 Zflops для сквозного обучения модели с триллионом параметров (т. е. 5, за которыми следуют24 нуля). Для обучения такой модели требуется 4000 A100 для работы с вычислительной эффективностью 50% в течение примерно 100 дней.

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

1.2 взвешивание

Далее мы рассмотрим компромиссы между параллелизмом данных, параллелизмом моделей и параллелизмом конвейера.

1.2.1 Параллелизм данных

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

эффективность памяти: Данные передаются при копировании модели и оптимизатора между всеми Worker'ами, поэтому эффективность использования памяти невысока.
Вычислительная эффективность: объем вычислений, выполняемых каждым работником, остается постоянным по мере увеличения степени параллелизма. Параллелизм данных может масштабироваться почти линейно в небольшом масштабе. Однако, поскольку стоимость связи для уменьшения градиентов между работниками положительно связана с размером модели, эффективность вычислений ограничена, когда модель велика или пропускная способность связи низка. Градиентное накопление — это распространенная стратегия, используемая для амортизации затрат на связь, которая может увеличить размер пакета и использовать его локально.micro-batchВыполните несколько прямых и обратных распространений, а затем уменьшите градиент перед выполнением обновления оптимизатора, тем самым амортизируя стоимость связи.

1.2.2 Параллелизм модели

модельный параллелизмэто другой тип техники. Он может разделить различные слои модели между несколькими работниками. По своей природе модельно-параллельные вычисления и обмен данными различаются в зависимости от структуры модели и, следовательно, требуют значительного объема работы для реализации. DeepSpeed использует возможности NvidiaMegatron-LMдля создания крупномасштабных модельно-параллельных языковых моделей на основе Transformer. Модельный параллелизм сокращает использование памяти пропорционально количеству рабочих процессов, что является наиболее эффективным с точки зрения использования памяти из трех режимов параллелизма. Но цена наименее вычислительно эффективна.

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

1.2.3 Конвейерный параллелизм

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

эффективность памяти: сокращение объема памяти при конвейерном параллелизме пропорционально количеству стадий конвейера, что позволяет линейно масштабировать размер модели с количеством рабочих процессов. Однако конвейерный параллелизм не уменьшает объем активной памяти каждого уровня. Кроме того, каждый рабочий процесс должен хранить активацию для каждого микропакета, работающего одновременно. Это приводит к тому, что память активации для первой стадии конвейера примерно такая же, как общая память активации для одной партии mirco.
Вычислительная эффективность: Поскольку трафик конвейера пропорционален только значению активации каждого уровня на границе этапа, параллельный трафик конвейера является самым низким. Однако он не может расширяться бесконечно. Как и параллелизм модели, увеличение размера конвейера снижает объем вычислений на каждом этапе конвейера, что снижает соотношение вычислений к обмену данными. Для достижения хорошей вычислительной эффективности конвейерный параллелизм также требует идеальной балансировки вычислительной нагрузки на каждом этапе.

Кроме того, параллелизм конвейера создает пузыри в начале и в конце каждой партии из-за необходимости пополнения или опорожнения трубопровода. Обучение с 4-кратным или 8-кратным увеличением количества этапов конвейера (и размера пакета) для шагов накопления градиента обеспечивает масштабируемость 81% и 90% соответственно по сравнению с одним этапом конвейера.

1.3 Параллеризм памяти и вычислительной эффективности

Параллелизм данных, модель и трубопровод Каждый играет определенную роль в улучшении памяти и вычислительной эффективности, поэтому Deepspeed сочетает в себе эти три мощных технологиях для обучения моделей триллионов и масштабирования до тысяч графических процессоров. Симбиотический параллель этих трех одновременно обращается к двум фундаментальным проблемам обучения моделей триллионов-параметра:эффективность памятииВычислительная эффективность, что позволяет масштабировать обучение глубокому обучению далеко за пределы того, что возможно с каждой стратегией в отдельности. Так что DeepSpeed может уместить огромные модели в видеопамять без ущерба для скорости.

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

Вычислительная эффективность: Чтобы масштабировать количество рабочих сверх того, что может поддерживать параллелизм модели и конвейера без ущерба для вычислительной эффективности, мы используем параллелизм данных с поддержкой ZeroRO (ZeRO-DP). ZeRO-DP может не только еще больше повысить эффективность использования памяти за счет разделения количества состояний оптимизатора, но и масштабироваться до любого количества графических процессоров с минимальными затратами на связь, используя отношения сопоставления на основе топологии связи.

Рис. 1: 3D-параллельный пример 32 рабочих. Слои нейронной сети разделены на четыре этапа конвейера. Уровни внутри каждого этапа конвейера далее делятся между четырьмя параллельными рабочими моделями. Наконец, каждый этап конвейера реплицируется между двумя экземплярами с параллельными данными, и ZeRO разделяет состояние оптимизатора между этими двумя копиями с параллельными данными.

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

Оптимизация пропускной способности внутриузловой и межузловой связи: Model-parallel имеет самые большие накладные расходы на связь из трех стратегий, поэтому мы отдаем приоритет размещению model-parallel worker groups внутри узлов, чтобы воспользоваться преимуществами большей пропускной способности внутри узла. Здесь мы выполняем параллелизм моделей с тензорными срезами на базе NVIDIA Megatron-LM. Когда модельная параллельная группа не может заполнить всех рабочих в узле, мы выбираем размещение параллельной группы данных в том же узле. В противном случае они будут параллельными данными между узлами. Конвейерный параллелизм имеет наименьший объем связи, поэтому мы можем планировать этапы конвейера между узлами, не ограничиваясь пропускной способностью связи.
Увеличьте пропускную способность за счет параллельного обмена данными: количество градиентов, которые должна пройти каждая группа параллельных данных, уменьшается линейно с масштабом конвейера и параллелизмом модели, поэтому общий 3D-трафик меньше, чем при чистом использовании параллельных данных. Кроме того, каждая параллельная группа данных взаимодействует независимо внутри небольшого числа рабочих локально, а межгрупповое общение может быть распараллелено друг с другом. Таким образом, уменьшая объем обмена данными и повышая локальность и параллелизм, мы можем эффективно расширить эффективную пропускную способность параллельного обмена данными.

Рисунок 2: Сопоставление рабочих процессов на рисунке 1 с графическими процессорами в системе из восьми узлов с четырьмя графическими процессорами в каждом. GPU одного цвета находятся на одном узле.

1.4 Как трехмерный параллелизм использует преимущества каждого параллелизма

Модель с триллионом параметров можно масштабировать на 4096 графических процессорах NVIDIA A100 с использованием 8-стороннего параллелизма модели, 64-стороннего параллелизма конвейера и 8-стороннего параллелизма данных.

Комбинируя модельный параллелизм и конвейерный параллелизм, трехмерный параллелизм обеспечивает превосходную эффективность использования памяти и высокую вычислительную эффективность на нескольких узлах. Параллелизм модели повышает эффективность активации и хранения состояния модели в узлах, а параллелизм конвейера (по сравнению с использованием только параллелизма модели) позволяет эффективно хранить состояние модели между узлами без ущерба для вычислительной эффективности.
Комбинируя модельный параллелизм с конвейерным параллелизмом, конвейерный параллелизм обеспечивает высокую эффективность вычислений с минимальными издержками даже при очень малых размерах пакетов. При параллелизме 8-процессорной модели каждая модель использует 1 микропакет, в результате чего эффективный размер микропакета составляет 1/8 на GPU. Таким образом, конвейерный параллелизм может достигать 90% вычислительной эффективности с использованием шагов накопления градиента, в 8 раз превышающих степень конвейерного параллелизма, а общий накопленный размер пакета на GPU составляет всего 1. В сочетании с параллелизмом данных это дает общий эффективный размер пакета 4096 на 4096 графических процессорах и по-прежнему обеспечивает эффективность конвейера 90%.

Но какую вычислительную эффективность обеспечивает параллелизм данных? Требует ли параллелизм данных больших пакетов на GPU для эффективности?

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

0x02 Введение

Мы следующиеZeRO: Memory Optimizations Toward Training Trillion Parameter Modelsучиться.

2.1 Резюме оригинального текста

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

Авторы статьи разработали новое решение, использующее оптимизатор с нулевой избыточностью (Zero) для оптимизации памяти, что значительно повышает скорость обучения при одновременном увеличении размера модели, которую можно эффективно обучить. ZeRO ищет подходящую среднюю точку между параллелизмом данных и параллелизмом моделей, которая надеется устранить избыточность памяти в данных и параллельном обучении моделей, сохраняя при этом низкий объем обмена данными и высокую степень детализации вычислений, что позволяет нам достичь устойчиво высокой эффективности. Эффективность масштабирует размер модели пропорционально количество устройств. Таким образом, ZeRO может получить преимущества параллелизма данных и параллелизма моделей. Та же память может использоваться для запуска более крупных моделей, а параллелизм данных можно использовать для обучения моделей, которые ранее можно было обучить только с помощью параллелизма моделей.

2.2 Введение в исходный текст

Общий параллелизм данных (DP) не уменьшает объем памяти на устройство, в то время как другие существующие решения, такие как конвейерный параллелизм (PP), модельный параллелизм (MP), разгрузка ЦП (CPU-Offloading) и т. д. Компромисс между функциональностью, удобством использования , а также память и вычислительная/коммуникационная эффективность. Среди различных существующих решений для обучения больших моделей MP может быть наиболее перспективным, однако MP не может масштабироваться до больших размеров. MP разделит модель по вертикали, разделив вычисления и параметры на каждом уровне между несколькими устройствами, что требует интенсивного обмена данными между каждым уровнем. Таким образом, они хорошо работают в пределах одного узла с высокой пропускной способностью связи между графическими процессорами, но за пределами одного узла эффективность быстро падает и, следовательно, не может эффективно масштабироваться за пределами одного узла.

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

Для больших моделей большую часть памяти занимает состояние модели, которое включает в себя состояния оптимизатора (такие как импульс и дисперсия в Адаме), градиенты и параметры.
Оставшуюся память занимают: активы, временные буферы и недоступные фрагменты, которые мы в совокупности называем оставшимся состоянием.

Поэтому мы разработали Zero Redundancy Optimizer, который может оптимизировать эффективность использования памяти в обоих аспектах, обеспечивая при этом высокую вычислительную и коммуникационную эффективность.

2.2.1 Оптимизация состояния модели

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

DP имеет хорошую эффективность вычислений/коммуникации, но низкую эффективность памяти, в то время как MP имеет низкую эффективность вычислений/коммуникации. В частности, DP реплицирует все состояние модели во всех параллельных процессах данных, что приводит к избыточному потреблению памяти; в то время как MP разделяет эти состояния для повышения эффективности использования памяти, это имеет тенденцию приводить к слишком мелким вычислениям и дорогостоящему обмену данными, тем самым уменьшая расширение. эффективность. Кроме того, эти методы статически поддерживают все состояния модели, необходимые на протяжении всего процесса обучения, даже если не все состояния модели всегда необходимы во время обучения.

Основываясь на этих наблюдениях, мы разрабатываем ZeRO-DP, который поддерживает вычислительную/коммуникационную эффективность DP при достижении эффективности памяти MP. ZeRO-DP устраняет избыточность состояния памяти в процессах, параллельных данным, за счет разделения состояния модели, а не его дублирования, так что потребление памяти на каждом графическом процессоре будет обратно пропорционально степени параллелизма данных, а также за счет использования динамической связи во время обучения. сохраняйте детализацию вычислений и объем связи в основном в соответствии с DP, чтобы можно было поддерживать эффективность вычислений/коммуникации.

Во время обучения модели большая часть памяти потребляется одной из трех вещей:

активация.
Состояние OGP, т.е. тензор, состоящий из состояния оптимизатора (O), градиента параметра (G) и самого параметра (P).
временный буфер.

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

ZeRO DP имеет три основных этапа оптимизации (показаны на рисунке 1 ниже), которые соответствуют разделам состояний оптимизатора, градиентов и параметров. При постепенном включении:

1) Optimizer State Partitioning (Pos): в 4 раза меньше памяти, столько же трафика, как и DP, этот этап также называется ZeRO-OS.

2) Добавлена оптимизация градиентного раздела (Pos+g): память уменьшена в 8 раз, а объем связи такой же, как у DP;

3) Добавлена оптимизация раздела параметров (Pos+g+p): уменьшение памяти линейно с параллелизмом DP. Память модели равномерно распределяется по каждому GPU, и потребление памяти на каждом GPU обратно пропорционально степени параллелизма данных, но трафик увеличивается лишь умеренно. Например, разделение на 64 графических процессора (Nd=64) приведет к сокращению памяти в 64 раза. Скромное увеличение трафика на 50%.

Потребление памяти можно увидеть на следующем рисунке:

Рис. 1. Сравнение потребления памяти на устройство на трех этапах оптимизации ZeRO-DP. ψ обозначает размер модели (количество параметров), K обозначает множитель памяти состояния оптимизатора, а Nd обозначает параллелизм DP, то есть Nd графических процессоров. В этом примере мы предполагаем обучение смешанной точности на основе оптимизатора Адама с размером модели ψ=7,5B, DP Nd=64 и K=12.

2.2.2 Оптимизация памяти остаточного состояния

После использования ZeRO-DP для оптимизации памяти, соответствующей состоянию модели, остаточная память (Residual State Memory) становится узким местом вторичной памяти, а оставшаяся память включает в себя: активацию, временный буфер и недоступные сегменты памяти. Мы разработали ZeRO-R для оптимизации оставшейся памяти, потребляемой этими тремя факторами по отдельности.

Для активаций (сохраненных из результатов прямого прохода для поддержки обратного прохода) мы заметили, что оптимизация контрольных точек могла бы помочь, но недостаточно для больших моделей. Таким образом, ZeRO-R оптимизирует активационную память, идентифицируя и удаляя копии активации в существующих схемах MP. Он также может разгрузить активацию ЦП, когда это необходимо.
ZeRO-R определяет подходящий размер временных буферов для достижения баланса памяти и вычислительной эффективности.
Мы наблюдаем, что во время обучения происходит некоторая фрагментация памяти из-за изменения времени жизни разных тензоров. Существование этих фрагментов может привести к сбою выделения памяти из-за отсутствия непрерывной памяти, даже если памяти достаточно. ZeRO-R активно управляет памятью в соответствии с различными жизненными циклами тензоров, чтобы предотвратить фрагментацию памяти.

ZeRO-DP и ZeRO-R объединяются, чтобы сформировать мощную оптимизированную для памяти систему для обучения DL, которую мы все вместе называем ZeRO.

2.2.3 Зеро и МП

Поскольку ZeRO устраняет неэффективность памяти в DP, естественно задаться вопросом: а нужен ли нам еще MP? когда вам это нужно? Как ZeRO работает с MP?

С ZeRO MP менее привлекателен для больших моделей. ZeRO-DP, по крайней мере, так же эффективен, как MP, в сокращении объема памяти для каждого устройства или более эффективен, когда MP не может разделить модель равномерно. Он также имеет сравнимую или лучшую эффективность масштабирования. Кроме того, параллелизм данных очень прост в использовании, поэтому он широко применим к различным рабочим нагрузкам, в то время как сегодняшние методы MP обычно требуют от разработчиков моделей дополнительной работы по модификации своих моделей, например, существующей работы (например, Megatron-LM), которая поддерживает только Ограниченный набор операций и моделей.

Тем не менее, есть ситуации, когда мы хотели бы воспользоваться MP:

i) MP может уменьшить объем памяти, занимаемой активацией для больших моделей при использовании с ZeRO-R.

ii) Для небольших моделей, где память активации не является проблемой. Когда используется только DP, хорошая сходимость может быть не достигнута, потому что совокупный размер партии слишком велик, и MP в это время также полезен. В этом случае ZeRO можно комбинировать с MP, чтобы получить модель приемлемого совокупного размера партии.

0x03 Связанная работа

3.1 Данные, модель и конвейерный параллелизм

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

Когда модель не помещается в памяти устройства, модельный параллелизм (MP) и конвейерный параллелизм (PP) разделяют модель между процессами по вертикали и по горизонтали соответственно.

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

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

PipeDream — еще одна реализация PP, которая хранит несколько копий устаревших параметров, чтобы скрыть пузыри конвейера без значительного увеличения размера пакета, что может снизить эффективность использования памяти. Кроме того, эта реализация не эквивалентна стандартному обучению DL и имеет значение для конвергенции обучения.

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

3.2 Непараллельная работа

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

Помимо MP и PP, есть много работы, направленной на снижение накладных расходов памяти при обучении DL.

3.2.1 Уменьшить память активации

В настоящее время проводится много работы, направленной на уменьшение объема памяти, занимаемой активацией, включая сжатие, контрольные точки активации или анализ в реальном времени. Эти усилия дополняют друг друга и работают вместе с ZeRO. На самом деле уменьшение памяти активации в ZeRO-R работает полностью параллельно с контрольной точкой активации.

3.2.2 CPU Offload

Есть также некоторые работы, использующие неоднородность вычислительных узлов для передачи состояния модели в память ЦП посредством разработки алгоритма или виртуализированной памяти, соответственно. Но это приводит к тому, что 50% времени тратится на передачу GPU-CPU-GPU. ZeRO отличается тем, что значительно снижает потребление памяти без необходимости сохранять состояние модели в памяти ЦП. В редких случаях ZeRO-R может разгрузить контрольную точку активации только для очень больших моделей, чтобы повысить производительность.

3.2.3 Оптимизатор эффективного использования памяти

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

3.3 Обучение оптимизатора

Для больших моделей методы адаптивной оптимизации (Adaptive) имеют решающее значение для достижения производительности и точности SOTA. По сравнению с SGD он поддерживает детализированную статистику первого и второго порядка для каждого параметра модели и градиента за счет значительного объема памяти. ZeRO может на порядки уменьшить объем памяти, занимаемой этими оптимизаторами, что делает эти сложные методы оптимизации практичными для обучения больших моделей на оборудовании со скромной памятью устройства. Это также позволяет разрабатывать и использовать более сложные оптимизаторы с большим объемом памяти и более конвергентные.

0x04 Куда девается память модели?

Давайте сделаем шаг назад и изучим потребление памяти текущей системой обучения. Например, модель GPT-2 с параметрами 1,5 млрд требует 3 ГБ памяти для весов (или параметров) с 16-битной точностью, но нельзя использовать Tensorflow или PyTorch для обучения на графическом процессоре с 32 ГБ памяти. Можно задаться вопросом, куда уходит вся память. Во время обучения модели большая часть памяти потребляется состояниями модели, то есть тензорами, состоящими из состояний оптимизатора, градиентов и параметров. В дополнение к этим состояниям модели остальная часть памяти потребляется активациями, временными буферами и фрагментированной памятью, которые мы называем остаточными состояниями. Мы подробно рассмотрим потребление памяти с этих двух аспектов.

4.1 Состояние модели: состояние оптимизатора, градиенты и параметры

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

Большая часть памяти устройства потребляется состоянием модели во время обучения. Например, возьмем Adam, одного из самых популярных оптимизаторов для обучения DL. Адаму необходимо хранить два состояния оптимизатора: i) усредненный по времени импульс и ii) дисперсию градиентов для вычисления обновлений. Следовательно, для обучения модели с помощью ADAM должно быть достаточно памяти для хранения копии импульса и дисперсии градиента. Кроме того, должно быть достаточно памяти для хранения градиентов и самих весов. Из трех типов тензоров, зависящих от параметров, состояние оптимизатора обычно потребляет больше всего памяти, особенно при применении обучения со смешанной точностью.

4.1.1 Обучение смешанной точности

Современный подход к обучению больших моделей на графических процессорах NVIDIA текущего поколения заключается в обучении со смешанной точностью (fp16/32), при котором параметры и активации сохраняются как fp16, что позволяет использовать блоки Tensor Core с высокой пропускной способностью на эти графические процессоры. Во время обучения со смешанной точностью оптимизатор выполняет прямое и обратное распространение с использованием весов и активаций fp16. Однако для эффективного вычисления и применения обновлений весов в конце обратного распространения оптимизатор смешанной точности должен хранить копии fp32 параметров, а также копии всех других состояний оптимизатора fp32.

Возьмем в качестве примера оптимизатор Adam. Для обучения модели смешанной точности с параметрами ψ с использованием Адама требуется достаточно памяти для хранения копий параметров и градиентов fp16. Его требования к памяти составляют 2ψ и 2ψ байта соответственно. Кроме того, необходимо сохранить fp32-копию состояния оптимизатора, импульса параметра и дисперсии с требованиями к памяти 4ψ, 4ψ и 4ψ байт соответственно.

Обозначим через K множитель памяти состояний оптимизатора, т.е. дополнительная память, необходимая для их хранения, составляет Kψ байт. Смешанная точность Адам имеет K = 12. В сумме это приведет к потребности в памяти 2ψ+2ψ+Kψ=16ψ байт. Для такой модели, как GPT-2 с 1,5 миллиардами параметров, требуется не менее 24 ГБ памяти, что намного больше, чем 3 ГБ, необходимых для сохранения одних только параметров fp16.

4.2 Использование оставшейся памяти

Первоначальное название было «Потребление остаточной памяти».

4.2.1 Активация

Во время обучения активации потребляют много памяти. В качестве конкретного примера модель GPT-2 с параметрами 1,5 млрд, обученная с длиной последовательности 1 тыс. и размером пакета 32, требует около 60 ГБ памяти. Контрольная точка активации (или пересчет активации) — это распространенный метод, который уменьшает память активации до квадратного корня из общего числа активаций при 33-процентных накладных расходах на пересчет. Это уменьшит потребление памяти активации для этой модели примерно до 8 ГБ.

Несмотря на значительное сокращение, память активации может стать довольно большой для более крупных моделей, даже с контрольными точками активации. Например, GPT-подобная модель со 100 миллиардами параметров потребует около 60 ГБ памяти для пакета размером 32, даже с контрольной точкой активации.

4.2.2 Временный буфер

Для больших моделей временные буферы, используемые для хранения промежуточных результатов, могут занимать много памяти. Некоторые операции, такие как уменьшение всех градиентов или вычисление нормы градиента, имеют тенденцию выполнять унифицированную операцию, объединяя все градиенты в один плоский буфер, что может повысить пропускную способность. Например, пропускная способность всех устройств уменьшается по мере увеличения объема сообщения. В то время как сами градиенты обычно хранятся как тензоры fp16, буферы слияния могут быть тензорами fp32 (в зависимости от типа операции). Эти временные размеры буфера очень важны, когда модель большая. Например, для модели с параметром 1,5Б плоский буфер fp32 требует 6ГБ памяти.

4.2.3 Фрагментация памяти

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

0x05 ZeRO: информация и обзор

ZeRO имеет два набора оптимизаций: i) ZeRO DP направлен на уменьшение объема памяти, занимаемой состояниями модели, ii) ZeRO-R направлен на сокращение оставшегося потребления памяти.

5.1 Идеи и обзор: ZeRO-DP

DP на базе ZeRO основан на трех ключевых идеях:

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

Основываясь на этих выводах, ZeRO DP достигает эффективности памяти MP, сохраняя при этом эффективность обучения DP. ZeRO DP разделяет состояния модели, а не копирует их, и использует план динамической связи, который использует присущие состояниям модели временные свойства при минимизации объема связи. Таким образом, ZeRO-DP линейно уменьшает объем памяти, занимаемой моделью для каждого устройства, по мере увеличения степени DP, сохраняя при этом трафик близким к DP по умолчанию, тем самым сохраняя эффективность.

5.2 Идеи и обзор: ZeRO-R

5.2.1 Уменьшение памяти активации

Два ключевых вывода:

MP разделяет состояние модели, но обычно требует копировать память активации. Например, если мы вертикально разделим параметры линейного слоя и вычислим их параллельно на двух графических процессорах, каждому графическому процессору потребуется полная активация для вычисления своих разделов.
Для моделей GPT-2 или более крупных моделей арифметическая интенсивность (отношение вычислений на итерацию к количеству контрольных точек активации на итерацию) очень велика (≥ 10 000) и увеличивается линейно со скрытым измерением, что позволяет скрыть стоимость перемещения данных для активации контрольных точек даже при низких пропускная способность.

ZeRO устраняет избыточность памяти в MP, разделяя контрольные точки активации между графическими процессорами и перестраивая их по требованию с помощью allgather. Уменьшение активационной памяти пропорционально степени МП. Для очень больших моделей ZeRO может даже выбрать перемещение разделов активации в память ЦП, сохраняя при этом хорошую эффективность из-за интенсивности вычислений в этих моделях.

5.2.2 Управление временными буферами

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

5.2.3 Управление фрагментацией памяти

Фрагментация памяти является результатом чередующегося выделения кратковременных и долгоживущих объектов памяти. Во время прямого распространения контрольные точки активации живут долго, а повторно вычисленные активации недолговечны. Кроме того, при обратном вычислении градиенты активации недолговечны, а градиенты параметров долгоживущие. Основываясь на этой реализации, ZeRO выполняет динамическую дефрагментацию памяти, перемещая контрольные точки активации и градиенты в предварительно выделенный непрерывный буфер памяти. Это не только повышает доступность памяти, но и повышает эффективность за счет сокращения времени, которое требуется распределителю памяти для поиска доступной непрерывной памяти.

0x06 Глубокое погружение в ZeRO-DP

В то время как существующие методы DP дублируют состояние модели на каждом устройстве и приводят к значительным накладным расходам памяти, ZeRO DP устраняет эту избыточность памяти, разделяя их (состояния оптимизатора, градиенты и параметры) между процессами, параллельными данным. На рис. 1 количественно и наглядно представлены требования к памяти с ZeRO-DP и без него. На рисунке показан объем памяти (1) состояния оптимизатора (2) градиента и (3) накопленной избыточности параметров после разделения. Мы называем это тремя этапами оптимизации ZeRO DP: Pos, Pg и Pp, которые мы подробно опишем ниже. Вот снова картинка.

6.1 Поз.: Разделение состояния оптимизатора

для $N_d$ Для DP параллелизма мы группируем состояния оптимизатора в $N_d$ В равном разделе такой i-й параллельный процесс данных обновляет только статус оптимизатора, соответствующий i-му разделу. Следовательно, каждому параллельному процессу данных необходимо только хранить и обновлять общий статус оптимизатора. $\frac{1}{N_d}$ , то только обновить $\frac{1}{N_d}$ параметры. В конце каждого шага обучения мы выполняем операцию сбора данных во всех параллельных процессах данных, чтобы получить полностью обновленные параметры во всех параллельных данных процессах.

Конкретный пример, показанный на рисунке 1, модель с B-параметрами 7,5, использует 64-канальный DP ( $N_d$ =64), для его Pos требуется 31,4 ГБ памяти. Для использования стандартного DP требуется 120 ГБ памяти. Кроме того, когда $N_d$ При увеличении требования к памяти для состояния модели уменьшаются с 4ψ+12ψ=16ψ байт до 4ψ+ $\ гидроразрыва {12ψ} {N_d}$ байт ≈ 4ψ байт, что приводит к сокращению в 4 раза.

6.2 Pg: Градиентное разбиение

Поскольку каждый параллельный процесс данных отвечает только за обновление соответствующего раздела параметров, каждый узел уменьшает градиент только той части параметров, за которую он отвечает. После слияния каждому узлу нужен только градиент, соответствующий его собственному разделу параметров, и он больше не нужен для других градиентов, поэтому их память может быть освобождена. Это уменьшает объем памяти градиента с 2ψ байт до $\ гидроразрыва {2ψ} {N_d}$ .

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

Экономия памяти: устраняя градиенты и избыточность состояний оптимизатора, мы дополнительно уменьшаем объем памяти до 2ψ+. $\ гидроразрыва {14ψ} {N_d}$ ≈ 2Ψ Как показано в примере на рисунке 1, для модели с B-параметром 7,5 требуется всего 16,6 ГБ памяти при использовании Pos+g и 64-канального DP (Nd=64) по сравнению со 120 ГБ при использовании стандартного DP. когда $N_d$ При увеличении требования к памяти для состояния модели уменьшаются с 2ψ+14ψ=16ψ байт до 2ψ+ $\ гидроразрыва {14ψ} {N_d}$ байт ≈ 2ψ байт, сокращение в 8 раз.

6.3 Pp: Раздел параметров

Так же, как состояния оптимизатора и градиенты, каждый процесс хранит только параметры, соответствующие его разделам. Когда для прямого и обратного распространения требуются параметры за пределами их разделов, они получаются от соответствующего параллельного процесса данных посредством широковещательной операции. Хотя на первый взгляд это может привести к значительным затратам на связь, мы обнаружили, что этот подход увеличивает общий объем связи базовой системы DP только в 1,5 раза, достигая при этом сокращения памяти, пропорционального Nd.

Экономия памяти: благодаря разделению параметров мы уменьшаем объем памяти, занимаемой параметрами ψ, с 16ψ до $\ гидроразрыва {16ψ} {N_d}$ . Как показано в примере на рис. 1, модель 7.5 В-параметров использует $P_{os+g+p}$ А для 64-канального DP (Nd=64) требуется всего 1,9 ГБ памяти, а для использования стандартного DP требуется 120 ГБ памяти.

Это имеет серьезные последствия: ZeRO-DP может работать с моделями любого размера, если имеется достаточное количество устройств для совместного использования состояния модели.

6.4 Влияние на размер модели

Три этапа разбиения Pos, Pos+g и Pos+g+p уменьшают потребление памяти каждым параллельным процессом данных в состоянии модели в 4 раза, 8 раз и 8 раз соответственно. $N_d$ раз. В таблице 1 анализируется потребление памяти состояния модели ZeRO-DP при 3 этапах ZeRO-DP на разных уровнях DP для нескольких примеров моделей.

Без ZeRO потребление памяти равно первой строке в таблице вне зависимости от степени DP. Обратите внимание, что когда Nd=64, ZeRO может использовать Pos, Pos+g и Pos+g+p для обучения моделей с параметрами до 7,5B, 14B и 128B соответственно. Когда Nd=1024, ZeRO со всеми включенными оптимизациями (Pos+g+p) может обучить модель с 1 триллионом параметров! А может модель любого размера! Без ZeRO самая большая модель, которую может запустить DP, имеет менее 1,5 миллиарда параметров.

0x07 Погрузитесь в ZeRO-R

7.1 $P_a$ : Контрольная точка активации раздела

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

В частности, после того как прямой проход для слоя в модели вычислен, входная активация распределяется по всем параллельным процессам модели до тех пор, пока она снова не понадобится на обратном проходе. На этом этапе ZeRO повторно материализует активную реплику, используя операцию сбора всех данных. Мы называем эту оптимизацию Pa. Он работает с активными контрольными точками и сохраняет только активные контрольные точки разделов, а не реплики. Кроме того, в случае очень больших моделей и очень ограниченной памяти устройства, контрольные точки активации для этих разделов также могут быть выгружены на ЦП, уменьшая накладные расходы памяти активации почти до нуля за счет дополнительных затрат на связь, которые мы называем $P_{a+cpu}$ .

Разделяя контрольные точки активации, ZeRO уменьшает площадь активации на коэффициент, пропорциональный степени MP. Рассмотрим обучение модели 100B с размером партии 32, длиной последовательности 1024 и степенью MP 16. Если мы проверяем одну активацию на слой преобразователя, то простое хранение контрольных точек активации потребует около 33 ГБ памяти на каждый GPU. Но если Pa равно нулю, емкость каждого графического процессора может быть уменьшена примерно до 2 ГБ. Кроме того, эти 2 ГБ могут быть выгружены на центральный процессор, что сокращает объем активной памяти почти до нуля.

7.2 CB: буфер фиксированного размера

ZeRO тщательно выбирает размер временного буфера данных, чтобы сбалансировать память и вычислительную эффективность. Во время обучения вычислительная эффективность некоторых операций может сильно зависеть от размера входных данных, причем большие входные данные более эффективны. Например, большая операция all-reduce получает более высокую пропускную способность, чем небольшая операция. Поэтому для повышения эффективности высокопроизводительные библиотеки, такие как NVIDIA Apex или Megatron, объединяют все параметры в один буфер перед выполнением этих операций. Однако накладные расходы памяти буфера слияния пропорциональны размеру модели. Например, для модели параметров 3B 32-битный буфер слияния потребует 12 ГБ памяти. Чтобы решить эту проблему, мы просто используем высокопроизводительные буферы постоянного размера, когда модель становится слишком большой. При этом размер буфера не зависит от размера модели, и, сохраняя размер буфера достаточно большим, мы все же можем добиться хорошей эффективности.

7.3 MD: Дефрагментация памяти

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

Ограниченная фрагментация памяти обычно не является проблемой при наличии достаточного объема памяти, но при обучении больших моделей с ограниченным объемом памяти фрагментация памяти может вызвать две проблемы: i) из-за отсутствия непрерывной памяти даже при достаточном количестве свободной памяти это может привести к OOM; ii) Неэффективность из-за того, что распределитель памяти тратит много времени на поиск смежных блоков памяти для удовлетворения запросов памяти.

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

0x08 Анализ трафика ZeRO-DP

Поскольку Zero улучшает размер модели за счет устранения избыточности памяти, естественно возникает вопрос, не в обмен ли это на эффективность использования памяти. Другими словами, каков трафик Zero-Powered для метода DP по сравнению с базовым методом DP? Ответ делится на две части: i) Zero-DP не генерирует дополнительной связи при использовании POS и PG, при этом достигается 8-кратное сокращение памяти; II) Zero-DP При использовании POS и PG До 1,5 раз больше связи, но при в то же время дальнейшее сокращение памяти в ND раз.

Параллельный трафик данных 8.1

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

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

8.2 Трафик ZeRO-DP

8.2.1 Трафик с использованием Pos+g

При градиентном разделении каждый процесс хранит только часть градиента, необходимую для обновления соответствующего раздела параметров. Следовательно, в отличие от all-reduce, ZeRO нужно выполнять только операции scatter-reduce на градиентах, что приводит к трафику ψ. После того как каждый процесс обновляет раздел для параметров, за которые он отвечает, выполняется сбор всех обновленных параметров из всех параллельных процессов данных. Это также приводит к трафику ψ. Таким образом, общий трафик для каждого шага обучения составляет ψ+ψ=2ψ, что в точности совпадает с базовым DP.

8.2.2 Трафик с использованием Pos+g+p

После разделения параметров каждый параллельный к данным процесс сохраняет только те параметры, за обновление которых он отвечает. Следовательно, при прямом распространении ему необходимо получить параметры всех остальных разделов. Однако это позволяет избежать накладных расходов на память за счет конвейерной обработки. Прежде чем вычислять прямое распространение для части модели, соответствующей конкретному разделу, процесс с параллельными данными, ответственный за этот раздел, может передавать веса всем процессам с параллельными данными. Эти параметры можно отбросить после завершения прямого распространения для этого раздела. Таким образом, общий трафик $\ гидроразрыва {ψ × N_d} {N_d} = ψ$ . Другими словами, мы полагаемся на распространение по всему прямому проходу, чтобы переупорядочить весь набор параметров и отбросить параметры после их использования. Обратите внимание, однако, что для обратного распространения эту сборку необходимо выполнить снова (но в обратном порядке).

Таким образом, общий трафик представляет собой сумму трафика, произведенного сокращениями и сбором, а общий объем равен 3ψ, что в 1,5 раза превышает базовый уровень. Как градиентное, так и параметрическое разбиение используют понимание того, что не все состояния градиента и параметра необходимы постоянно, а вместо этого оптимизируют память, разумно передавая состояния.

0x09 Анализ связи ZeRO-R

Мы сравнили трафик активированных секциями контрольных точек (Pa) в ZeRO-R с базовым MP и показали, что увеличение трафика, вызванное Pa, обычно составляет менее одной десятой от базового MP. Кроме того, мы анализируем коммуникационные издержки трафика Pa по сравнению с трафиком DP, чтобы выявить случаи, когда Pa повышает эффективность, позволяя использовать большие пакеты и уменьшая трафик DP. Мы используем этот анализ, чтобы решить, следует ли и когда применять Pa и Pa+cpu.

Компромисс трафика для активируемых разделом контрольных точек зависит от размера модели, стратегии контрольных точек и стратегии MP. Чтобы поделиться конкретными идеями, мы проводим наш анализ в контексте модели, реализованной с использованием метода SOTA MP (Megatron-LM).

В Megatron-LM с контрольной точкой активации каждый преобразователь выполняет две операции all-reduce размером batch×seq_length×hidden_dim при прямом распространении, две операции all-reduce используются для прямого пересчета, две другие операции all-reduce используются для обратного распространения. Общий трафик для каждого блока равен 12 × длина последовательности × скрытый размер, поскольку трафик для всех редукций равен 2 × размеру сообщения.

Когда ZeRO-R разделяет контрольные точки активации, необходимо выполнить дополнительную операцию сбора, прежде чем обратное распространение будет пересчитано вперед для каждой контрольной точки активации. Как правило, мы проверяем входные активации каждого блока преобразователя, и каждый блок преобразователя требует сбора всех данных. Таким образом, коммуникационные издержки Pa равны длине последовательности ∗ скрытому размеру, поскольку весь агрегированный трафик представляет собой размер сообщения. Таким образом, общие коммуникационные издержки Pa составляют менее 10 % исходного объема коммуникационных данных параллелизма модели.

Когда MP сочетается с DP, Pa можно использовать для уменьшения трафика, параллельного данным, на порядок, в то время как трафик, параллельный модели, увеличивается на 10% и значительно повышает эффективность, когда параллельная связь данных становится узким местом производительности. Обратите внимание, что Pa снижает потребление памяти активации до параллелизма MP, позволяя пропорционально увеличить размер пакета. Для больших моделей количество MP может достигать 16 (#GPU на узлах DGX-2), что позволяет увеличить размер пакета до 16 раз. Трафик для параллельного обучения данных обратно пропорционален размеру пакета. Таким образом, увеличение размера пакета на порядок из-за Pa может привести к уменьшению на порядок трафика, параллельного данным.

Наконец, если будет принят Pa+cpu, контрольная точка активации раздела будет перенесена на процессор, потребность в памяти для активации будет снижена почти до нуля, а по сравнению с Pa перемещение данных между памятью процессора увеличится в 2 раза. В крайних случаях трафик DP является основным узким местом, потому что даже с Pa размер пакета невелик, пока накладные расходы на передачу данных процессора меньше, чем накладные расходы на трафик DP, Pa+cpu может повысить эффективность за счет увеличения размера пакета. который обычно хорошо работает с небольшими партиями.

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

0xEE Личная информация

★★★★★★Думая о жизни и технологиях★★★★★★

Публичный аккаунт WeChat:мысли Росси

ссылка 0xFF

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

DeepSpeed: инструмент обучения гипермасштабируемой модели для всех

DeepSpeed: Extreme-scale model training for everyone