[Анализ исходного кода] Как Facebook обучает очень большую модель --- (1)

0x00 сводка

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

FSDP (Fully Sharded Data Parallel) — это обновленная версия PyTorch DDP, предложенная Facebook после подробного ознакомления с Microsoft ZeRO, которую можно рассматривать как эталон по сравнению с Microsoft ZeRO, и ее суть — разделение параметров. Разделение параметров заключается в разделении параметров модели на каждый графический процессор. Мы используем документы, блоги и код Google, Microsoft и Facebook для анализа обучения.

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

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

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

[Перевод статьи] Распределенное обучение Разделение параметров Google Weight Sharding

0x01 Введение

1.1 FAIR & FSDP

Обучать модели ИИ в масштабе непросто. Помимо того, что требуется много вычислительной мощности и ресурсов, обучение очень больших моделей связано со значительной инженерной сложностью. Facebook Artificial Intelligence Research (FAIR) Engineering работает над созданием инструментов и инфраструктуры, чтобы упростить обучение больших моделей ИИ.

Fully Sharded Data Parallel (FSDP) — новейший инструмент, представленный FAIR. Он разделяет параметры модели ИИ между параллельно работающими с данными работниками и при необходимости переносит часть обучающих вычислений на ЦП. Как следует из названия, FSDP — это алгоритм параллельного обучения данных. Хотя параметры распределяются по разным графическим процессорам, вычисление каждого микропакета остается локальным для каждого рабочего графического процессора. Эта концептуальная простота делает FSDP более простым для понимания и более применимым к различным сценариям использования (по сравнению с внутриуровневым параллелизмом и конвейерным параллелизмом). По сравнению с параллельным методом сегментирования данных состояния + градиента оптимизатора, сегменты FSDP распределяют параметры модели более равномерно благодаря обмену данными и перекрытию вычислений во время обучения, что приводит к повышению производительности.

FSDP может более эффективно обучать модели на несколько порядков, используя меньше графических процессоров. FSDP реализован в библиотеке FairScale, что позволяет инженерам и разработчикам масштабировать и оптимизировать обучение своих моделей с помощью простого API. В Facebook FSDP был интегрирован и протестирован для обучения некоторых моделей NLP и Vision.

1.2 Требования к вычислительной мощности для крупномасштабного обучения

Для обучения крупномасштабной модели требуется много вычислительных ресурсов, таких как OpenAI GPT-3 со 175 миллиардами параметров. По оценкам, на его обучение уйдет 355 лет работы графического процессора, что эквивалентно 1000 графическим процессорам, работающим непрерывно более 4 месяцев.

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

Напротив, FSDP не делает никаких компромиссов. Он повышает эффективность использования памяти за счет разделения параметров модели, градиентов и состояний оптимизатора на графическом процессоре, а также повышает эффективность вычислений за счет разложения связи и ее перекрытия с прямыми и обратными процессами. FSDP дает те же результаты, что и стандартное обучение распределенным параллельным данным (DDP), и предоставляет простой в использовании интерфейс, который является альтернативой модулю PyTorch Distributed Data Parallel. Ранние тесты в Facebook показывают, что FSDP может масштабироваться до триллионов параметров.

0x02 Как работает FSDP

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

2.1 Полное разделение параметров

Один из способов уменьшить количество реплик — применить процесс полного разделения параметров, при котором предоставляется только подмножество параметров модели, градиентов и оптимизаторов, необходимых для локальных вычислений. Реализация этого подхода, ZeRO-3, была популяризирована Microsoft.

Ключом к разблокировке полной сегментации параметров является то, что мы можем разложить все операции сокращения в DDP на отдельные операции сокращения-разброса и все-сборки.

Изображение из:Engineering.convenience.com/Боюсь-контент/…

«Все-уменьшить» — это комбинация «уменьшить-разбросать» и «все-собрать». Стандартную операцию «Все-уменьшить» для агрегации градиентов можно разбить на два отдельных этапа: «уменьшить-разбросать» и «все-собрать».

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

Благодаря перераспределению «уменьшить разброс» и «все собрать» каждому рабочему процессу DDP нужно хранить только один фрагмент параметра и состояние оптимизатора.

2.2 Сравнение

На рисунке ниже показано стандартное обучение DDP (верхняя половина) и обучение FSDP (нижняя половина):

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

• В ФСДП:

  • Model shard: существует только на каждом графическом процессореШардинг модели.
  • All-gather: каждый GPU собирает все данные с других GPU через all-gatherВеса, чтобы локально вычислить прямое распространение. Это подчеркнутая часть идеи диссертации, стр.
  • Вперед (местный): выполнить операцию переадресации локально. И прямое вычисление, и обратное вычисление используют полную модель.
  • All-gather: затем сделайте это снова перед обратным распространениемВесасобирать. Это подчеркнутая часть идеи диссертации, стр.
  • Назад (местный): локальное выполнение операции в обратном направлении. И прямые, и обратные вычисления используют полную модель, и каждый GPU такжевсе градиенты.
  • Reduce-scatter: после обратного распространения локальныйградиентОн агрегируется и сегментируется на каждом графическом процессоре с помощью редукции-разброса Градиент на каждом сегменте — это часть, соответствующая этому разделу после агрегации, что является подчеркнутой частью идеи статьи Pg.
  • Обновить вес (локальный): каждый GPU обновляет свой локальныйВесаФрагментация.

Чтобы максимизировать эффективность использования памяти, мы можем отбросить все веса после прямого прохода каждого слоя, сохраняя память для последующих слоев. Этого можно добиться, применив оболочку FSDP к каждому уровню в сети (установив reshard_after_forward=True).

Вот реализация псевдокода:

FSDP forward pass:
    for layer_i in layers:
        all-gather full weights for layer_i # 权重
        forward pass for layer_i
        discard full weights for layer_i # 权重

FSDP backward pass:
    for layer_i in layers:
        all-gather full weights for layer_i # 权重
        backward pass for layer_i
        discard full weights for layer_i # 权重
        reduce-scatter gradients for layer_i # 梯度

2.3 Кардочесание

Давайте посмотрим на FSDP в сочетании с идеями статьи.

2.3.1 Идеи

Идея дипломной работы заключается в следующем:

• Pp: Parameter Partitioning, каждый процесс хранит только параметры, соответствующие его разделу. когдаПрямое и обратное распространениеКогда требуются параметры за пределами своих разделов, они принимаются от соответствующего параллельного процесса данных посредством широковещательной операции. Хотя на первый взгляд это может привести к значительным затратам на связь, мы обнаружили, что этот подход увеличивает общий объем связи базовой системы DP только в 1,5 раза, при этом достигается сокращение объема памяти, пропорциональное Nd.
• Pos : Optimizer State Partitioning, для$N_d$Для DP параллелизма мы группируем состояния оптимизатора в$N_d$равные разделы, так что i-й процесс с параллельными данными обновляет только состояние оптимизатора, соответствующее i-му разделу. Таким образом, каждому параллельному процессу данных необходимо хранить и обновлять только общее состояние оптимизатора.$\frac{1}{N_d}$, то только обновить$\frac{1}{N_d}$параметры. В конце каждого шага обучения мы выполняем операцию сбора данных во всех параллельных процессах данных, чтобы получить полностью обновленные параметры во всех параллельных данных процессах.
• Pg: Gradient Partitioning, поскольку каждый процесс с параллельными данными отвечает только за обновление соответствующего раздела параметров, каждый узел отвечает только за градиент той части параметров, за которую он отвечает.Статут. После слияния каждому узлу нужен только градиент, соответствующий его собственному разделу параметров, и он больше не нужен для других градиентов, поэтому их память может быть освобождена. Это уменьшает объем памяти градиента с 2ψ байт до$\ гидроразрыва {2ψ} {N_d}$. На самом деле, этоОперация уменьшения разброса, градиенты разных параметров сводятся к разным процессам.

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

2.3.2 Этапы процесса

Давайте снова покажем конкретный процесс. Если предположить, что параллелизм данных равен n, то имеется n GPU, тогда 1/n от общих параметров модели сохраняется на каждом GPU, а состояние градиента и оптимизатора естественным образом разделяется, и на каждом GPU присутствует параллелизм данных.

• Исходное состояние: поверх каждого графического процессора$P_n, G_n, O_n$. Обратите внимание, что поскольку модель на этом графическом процессоре$P_n$,так$O_n$естественно соответствует$P_n$, он автоматически фрагментируется.
• При расчете вперед каждый$GPU_n$Все берут на себя ответственность за параметры$P_n$Рассылка на все остальные графические процессоры после прямого вычисления, каждый$GPU_n$Оба получают свои собственные входные обучающие данные$data_n$Потеря$loss_n$.
• В обратном расчете каждый$GPU_n$также поставить свои ответственные параметры$P_n$Трансляция на все остальные графические процессоры и, наконец, вычисление соответствующих данных.$data_n$градиент$G_n$.
• градиент$G_n$агрегировать в соответствующий$GPU_n$вверх, на этот раз$GPU_n$Градиент на$reduce(G_0, ..., G_n)$Часть, соответствующая вашему рангу. Обратите внимание, что в процессе агрегации градиента используется метод уменьшения разброса, потому что каждому графическому процессору нужно обновлять только свою часть.$P_n,G_n,O_n$, так что нет необходимости все собирать.

0x03 How to use FSDP

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

3.1 Использование FSDP в языковой модели

Для языковых моделей вы можете передать следующие новые параметры вfairseq frameworkFSDP поддерживается среди:

• --ddp-backend = oll_sharded: включить полностью осколки через FSDP.
• --cpu-offload: Разгрузить состояние оптимизатора и копию модели FP32 на процессор (в сочетании с --optimizer=cpu_adam).
• --no-reshard-after-forward: Увеличить скорость обучения для больших моделей (параметры 1B+), аналогично этапу 2 ZeRO.
• Другие распространенные опции (–fp16, –update-freq, –checkpoint-activations, –offload-activations и т. д.) продолжают нормально работать.

Для получения подробной информации см.учебник по ярмарке.

3.2 Использование FSDP в моделях компьютерного зрения

Для моделей компьютерного зрения,VISSLFSDP может поддерживаться в regnet и тестироваться на архитектуре regnet. Слои, такие как BatchNorm и ReLU, были беспрепятственно обработаны и протестированы на сходимость. FSDP можно включить с помощью опций ниже.

• config.MODEL.FSDP_CONFIG.AUTO_SETUP_FSDP=True
• config.MODEL.SYNC_BN_CONFIG.SYNC_BN_TYPE=pytorch
• config.MODEL.AMP_PARAMS.AMP_TYPE=pytorch

Продолжить исследование можно по следующей ссылкеthis section.

3.3 Использование FSDP в PyTorch Lightning

Для упрощения интеграции с более общими вариантами использования PyTorch Lightning сделала FSDP бета-функцией. [этот урок](py torch-lightning.read the doc S.IO/en/latest/ ах…training) содержит подробный пример использования подключаемого модуля FSDP с PyTorch Lightning. Его можно активировать, добавив plugins='fsdp', как показано ниже.

model = MyModel()
trainer = Trainer(gpus=4, plugins='fsdp', precision=16)
trainer.fit(model)

trainer.test()
trainer.predict()

3.4 Использование библиотеки FSDP непосредственно из FairScale

Основной библиотекой разработки для FSDP являетсяFairScale.. Вы можете напрямую использовать FSDP FairScale в следующем примере, просто заменив DDP.

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
...
# sharded_module = DDP(my_module)
sharded_module = FSDP(my_module)
optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
for sample, label in dataload.next_batch:
  out = sharded_module(x=sample, y=3, z=torch.Tensor([1]))
  loss = criterion(out, label)
  loss.backward()
  optim.step()

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

1. **Упаковка моделей.** Чтобы свести к минимуму краткосрочные потребности в памяти графического процессора, пользователям необходимо упаковывать модели вложенным образом. Это добавляет сложности, но очень полезно при переносе существующего кода модели PyTorch.
2. Инициализация модели:В отличие от DDP, FSDPНе будуАвтоматически синхронизируйте веса моделей между рабочими процессами графического процессора. Это означает, что инициализация модели должна выполняться осторожно, чтобы все рабочие процессы GPU имели одинаковые начальные веса.
3. **Настройки оптимизатора.** Из-за сегментирования и упаковки FSDP поддерживает только определенные типы оптимизаторов и их настройки. В частности, если модуль упакован в FSDP и его параметры объединены в один тензор, пользователь не может использовать разные гиперпараметры для разных групп параметров в таком модуле.
4. **Смешанная точность**: FSDP поддерживает расширенную тренировку со смешанной точностью с основным весом FP16, а также уменьшением и разбросом по типу FP16 на градиентах. Однако некоторые части модели могут сойтись только при использовании полной точности, в этих случаях требуется дополнительная упаковка для выборочного запуска некоторых частей модели с полной точностью.
5. **Контрольные точки состояния и логические выводы.** Если модель большая, сохранение и загрузка состояния модели может стать затруднительной. FSDP поддерживает несколько методов, позволяющих сделать эту задачу возможной, но за эти методы приходится платить.
6. Наконец, FSDP обычно ассоциируется сАктивировать контрольную точкуфункционируют вместе, напримерcheckpoint_wrapper. Пользователям может потребоваться тщательно настроить стратегию контрольных точек активации, чтобы разместить большую модель в ограниченном пространстве памяти графического процессора.

0x04 управление памятью

Далее мы рассмотрим, как FSDP управляет памятью.

FairScale предоставляетZeRO <https://arxiv.org/pdf/1910.02054.pdf>Вдохновленный алгоритм: при обучении с параллелизмом данных вам необходимо найти компромисс между использованием памяти с точки зрения эффективности вычислений/коммуникации. С другой стороны, при использовании параллельного обучения модели возникает компромисс между вычислениями и обменом данными для памяти.

Использование памяти для обучения модели обычно делится на две категории:

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

• Оставшееся состояние: активно, временный буфер, фрагментированная память.

Чтобы уменьшить избыточность состояний модели, ZeRO предлагает три различных алгоритма. Они реализованы в FairScale как Optimizer State Sharding (OSS), Sharded Data Parallel (SDP) и, наконец, Fully Sharded Data Parallel (FSDP). Давайте углубимся в реальную механику каждого алгоритма и поймем, почему они экономят память.

4.1 Optimizer State Sharding (OSS)

FairScale реализовала OSS, оптимизированную для памяти, связанную с памятью оптимизатора.

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

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

4.1.1 Процесс обучения

Процесс обучения можно изменить по сравнению с процессом выполнения DDP следующим образом:

1. обернутый оптимизатор разделяет состояние оптимизатора жадным алгоритмическим способом на основе размера параметра (а не порядка использования). Это делается для того, чтобы каждый ранг занимал почти одинаковый объем памяти оптимизатора.

2. Процесс обучения аналогичен процессу распределенного параллелизма данных (DDP) PyTorch. Передний проход выполняется на каждом ряду, за которым следует обратный проход. Во время обратного прохода используйте allreduce для синхронизации градиентов.

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

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

OSS полезен, когда вы используете оптимизатор с дополнительным состоянием (например, Adam). Если вы используете SGD или любой другой оптимизатор с ограниченным объемом памяти, вы можете заметить замедление при использовании нескольких узлов из-за дополнительной связи на шаге 4. Во время процесса allreduce на шаге 2 также тратится некоторая часть памяти для хранения градиентов, которая затем отбрасывается.

4.1.2 Лучшие практики

• OSS предоставляет флаг широковещания_fp16, который вам, вероятно, следует использовать в многоузловых заданиях. Это обычно не требуется в экспериментах с одним узлом.

• Если ваша модель крайне несбалансирована по размеру (например, имеется огромный тензор), то этот подход мало чем поможет, а варианты нарезки тензора типа «fairscale.nn.FullyShardedDataParallel» будут более желательны.

• 3. OSS должно быть временным решением в среде DDP, которое остается совместимым с большинством функций DDP.

4.1.3 Производительность

• На одном узле OSS всегда должен быть быстрее, чем vanilla PyTorch, экономия памяти будет варьироваться в зависимости от используемого оптимизатора.

• OSS также может быть быстрее или медленнее, чем ванильный PyTorch, при использовании нескольких узлов, в зависимости от используемого оптимизатора и дополнительных флагов (таких как Broadcast_fp16, сжатие градиента, накопление градиента, упомянутое выше).

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

4.2 Optimizer + Gradient State Sharding

Хотя OSS решает проблему избыточности в оптимизаторе, по-прежнему существует дублирование вычислений агрегации градиентов и дополнительная память для градиентов. Чтобы преодолеть избыточную память градиента, мы можем использовать сегментацию градиента илиZeRO-2. Это было реализовано API Sharded Data Parallel (SDP) в FairScale.

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

4.2.1 Процесс обучения

Процесс обучения выглядит следующим образом:

1. Как и прежде, обернутый оптимизатор разбивает параметры по разным группам столбцов.

2. Модель теперь обернута с помощью оболочки Sharded Data Parallel (SDP), которая позволяет нам добавлять соответствующие крючки и поддерживать состояние во время обучения.

3. SDP фокусируется на обучаемых параметрах и добавляет обратный хук для каждого параметра.

4. Во время обратного распространения градиент будет уменьшен до указанного ранга, который указан в 1 как часть процесса разделения. Используйте reduce op вместо allreduce op, тем самым уменьшая коммуникационные накладные расходы.

5. Каждый ранг обновляет параметры, за которые он отвечает.

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

И OSS, и SDPAPI позволяют уменьшить объем памяти, используемый для градиентов и состояния оптимизатора, но если сеть работает медленно, могут возникнуть дополнительные расходы на связь. При возникновении проблем с нехваткой памяти (OOM) вы можете в качестве первого шага попробовать OSS и SDP.

4.2.2 Передовой опыт

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

• Если на одном узле, как правило, лучше не использовать 'reduce_buffer_size', так как это влечет за собой затраты на задержку, но не увеличивает память. Установите это значение на 0, чтобы не использовать эту функцию.

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

4.3 Optimizer + Gradient + Horizontal Model Sharding

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

Нарезка параметров аналогична градиентам и состояниям оптимизатора, т. е. каждый ранг параллельных данных отвечает за один сегмент параметров модели. FairScale реализует сегментирование параметров с помощью API Fully Sharded Data Parallel (FSDP), который хорошо принятZeRO-3 <https://arxiv.org/pdf/1910.02054.pdf>вдохновение.

Разделение параметров имеет два ключевых момента:

• Все операции сокращения можно разделить на операции сокращения и сбора, как и в предыдущих методах сегментирования (состояния оптимизатора и градиенты).

• Отдельные слои можно обернуть с помощью FSDP API, что позволяет нам передавать все параметры, необходимые для одного слоя, на данный GPU в данном экземпляре, вычислять прямой проход, а затем отбрасывать параметры, которые не относятся к этому рангу.

Использовать FSDP так же просто, как заменить оригинальный DDP в коде. Примечание. В настоящее время FSDP требует, чтобы модельnn.SequentialМодель.

from torch.utils.data.dataloader import DataLoader
from torchvision.datasets import FakeData
from torchvision.transforms import ToTensor

from fairscale.experimental.nn.offload import OffloadModel

num_inputs = 8
num_outputs = 8
num_hidden =  4
num_layers =  2
batch_size =  8

transform = ToTensor()
dataloader = DataLoader(
    FakeData(
        image_size=(1, num_inputs, num_inputs),
        num_classes=num_outputs,
        transform=transform,
    ),
    batch_size=batch_size,
)

model = torch.nn.Sequential(
    torch.nn.Linear(num_inputs * num_inputs, num_hidden),
    *([torch.nn.Linear(num_hidden, num_hidden) for _ in range(num_layers)]),
    torch.nn.Linear(num_hidden, num_outputs),
)

4.3.1 Процесс обучения

Конкретный тренировочный процесс выглядит следующим образом:

• Прежде чем приступить к расчету конкретного слоя,allgatherПараметры, необходимые для прямого распространения каждого слоя модели.

• Расчет ведется вперед.

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

• Вычисление распространяется обратно.

• Уменьшите градиенты так, чтобы агрегированные градиенты накапливались по рангам, отвечающим за соответствующие параметры.

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

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

FSDP также поддерживает обучение со смешанной точностью, при котором вычисления и обмен данными выполняются с точностью FP16. Если вы хотите уменьшить то, что вы делаете в FP32 (что является поведением DDP по умолчанию), вы должны установитьfp32_reduce_scatter=True.

Для дополнительной экономии памяти FSDP поддерживает разгрузку неиспользуемых в настоящее время параметров и градиентов на ЦП. Это можно включить, установив для параметров «move_params_to_cpu» и «move_grads_to_cpu» значение True.

4.3.2 Передовой опыт

• Для ФСДП лучше использоватьmodel.zero_grad(set_to_none=True), так как это экономит много памяти после пошагового выполнения.

• torch.cuda.amp.autocastПолностью совместим с FSDP. Вам нужно установить для аргумента «mixed_precision» значение True.

• В сочетании с активацией контрольных точек лучше использовать FSDP(checkpoint_wrapper(module)) вместо checkpoint_wrapper(FSDP(module)). Последнее приведет к большему количеству связи и более низким скоростям.

• FSDP совместим с DDP, которые используют точечные оптимизаторы, такие как Adam, AdamW, ADADDelta, Adamax, SGD и т. д. Разделение приведет к немного другим результатам при использовании неточечных оптимизаторов (например, Adagrad, Adafactor, LAMB и т. д.).

4.3.3 Производительность

• Для лучшей эффективности памяти используйте «auto_wrap», чтобы обернуть каждый слой в сети с помощью FSDP и обернутьreshard_after_forwardУстановите значение «Истина». Это будет медленнее, но затраты памяти будут минимальными.

• Для оптимальной скорости тренировки используйтеreshard_after_forwardУстановите значение False (не нужно обертывать каждый слой, но если установлено, скорость еще больше увеличится).

Были представлены поддержка, основные принципы FSDP и способы его использования.В следующей статье мы представим детали его кода и посмотрим, как минимизировать использование памяти.

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

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

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

ссылка 0xFF

Fully Sharded Data Parallel: faster AI training with fewer GPUs

ZeRO и DeepSpeed: оптимизация для обучающих моделей с более чем 100 миллиардами параметров (Microsoft)

Fully Sharded Data Parallel: faster AI training with fewer GPUs

GitHub.com/Microsoft/D…

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training