ZeRO-1 vs ZeRO-2 vs ZeRO-3: что и когда использовать?

Краткий тезис

ZeRO (Zero Redundancy Optimizer) — это техника шардинга состояния оптимизатора, градиентов и параметров модели при распределённом обучении, реализованная в библиотеке DeepSpeed. Она позволяет обучать большие модели, которые не помещаются в память одной GPU, за счёт распределения данных между устройствами. Выбор стадии ZeRO зависит от размера модели, количества GPU и пропускной способности сети: ZeRO-1 шардит только optimizer states (экономия ~2x), ZeRO-2 добавляет шардинг градиентов (~4x), ZeRO-3 шардит ещё и веса модели (экономия до Nx, где N — число GPU). Для моделей до 13B параметров обычно достаточно ZeRO-2, для 30B+ рекомендуется ZeRO-3, но с учётом дополнительных коммуникационных накладных расходов.

1. Проблема памяти при обучении больших моделей

При обучении нейронной сети на GPU память расходуется на четыре основные компонента:

• Параметры модели (weights) — сами веса.
• Градиенты — производные loss по весам, нужны для обратного распространения.
• Состояние оптимизатора (optimizer states) — например, для Adam это моменты первого и второго порядка (m, v), которые занимают в 2–3 раза больше памяти, чем сами веса.
• Активации (activations) — промежуточные значения слоёв, сохраняемые для backward pass.

Для модели с 1 млрд параметров в FP16 (2 байта на параметр) веса занимают 2 ГБ, градиенты — ещё 2 ГБ, а состояние Adam (обычно FP32) — 8 ГБ (4 байта на параметр для m и v). Итого только для обучения нужно ~12 ГБ, плюс активации — легко превысить лимит одной GPU (например, 16 ГБ). ZeRO решает эту проблему, распределяя эти компоненты между несколькими GPU, а не дублируя их на каждой.

2. Общая идея ZeRO

В обычном Data Parallel|Data Parallelism (DP) каждая GPU хранит полную копию модели, градиентов и состояния оптимизатора. После forward/backward градиенты усредняются (AllReduce), и каждая GPU обновляет свою копию. Это приводит к огромному дублированию памяти.

ZeRO предлагает шардинг (разделение) этих компонентов между GPU, устраняя избыточность. Каждая GPU хранит только свою часть данных, а при необходимости (forward/backward) недостающие части собираются через коммуникационные операции (AllGather, ReduceScatter). ZeRO не меняет количество вычислений, но увеличивает объём коммуникаций.

3. ZeRO-1: шардинг состояния оптимизатора

Что шардится только optimizer states (например, m и v для Adam).
Экономия памяти примерно в 2 раза по сравнению с обычным DP (если состояние оптимизатора занимает больше всего памяти).
Коммуникация не добавляет новых операций, так как градиенты всё равно усредняются через AllReduce, а после обновления весов каждая GPU имеет полную копию весов.
Когда использовать для небольших моделей (до ~1B), где основное узкое место — состояние оптимизатора. На практике ZeRO-1 редко применяется отдельно, так как ZeRO-2 даёт больше выгоды с минимальными дополнительными накладными расходами.

4. ZeRO-2: шардинг градиентов

Что шардится optimizer states + градиенты.
Экономия памяти примерно в 4 раза по сравнению с DP (устраняется дублирование градиентов и состояний).
Коммуникация вместо AllReduce для градиентов используется ReduceScatter (суммирование и распределение частей градиентов по GPU), затем каждая GPU обновляет только свою часть состояния оптимизатора. После обновления веса остаются полными на каждой GPU (не шардятся).
Когда использовать стандартный выбор для моделей до 13B параметров. Обеспечивает хороший баланс между экономией памяти и коммуникационными затратами.

5. ZeRO-3: шардинг весов модели

Что шардится optimizer states + градиенты + параметры модели (weights).
Экономия памяти в N раз (где N — число GPU) по сравнению с DP. Например, на 8 GPU каждая хранит 1/8 весов, градиентов и состояний.
Коммуникация самая интенсивная. Перед каждым forward/backward слоя необходимо выполнить AllGather для сбора полных весов этого слоя на всех GPU. После вычислений веса снова шардятся (никуда не отправляются, просто хранятся части). Это приводит к значительному увеличению времени коммуникации, особенно при большом числе GPU или медленной сети.
Когда использовать для очень больших моделей (30B+), которые не помещаются в память одной GPU даже с ZeRO-2. Также полезен, когда нужно обучать модель на ограниченном числе GPU (например, 4–8) с большим размером.

6. Сравнительная таблица ZeRO-1 / ZeRO-2 / ZeRO-3

7. Когда что использовать: практические рекомендации

Выбор стадии ZeRO определяется тремя факторами: размер модели, количество GPU и пропускная способность сети (NVLink, InfiniBand или Ethernet).

• Модели до 1B параметров можно использовать обычный Data Parallelism или ZeRO-1, но проще включить ZeRO-2 — он даёт запас памяти без существенного замедления.
• Модели 1B–13B оптимален ZeRO-2. Экономит память в 4 раза, коммуникационные накладные расходы приемлемы. Если сеть медленная (Ethernet), можно попробовать ZeRO-1, но обычно ZeRO-2 предпочтительнее.
• Модели 13B–30B если количество GPU достаточно (например, 8–16), ZeRO-2 может работать, но уже на грани. Рекомендуется ZeRO-3 с включённой опцией offload_optimizer (выгрузка состояния оптимизатора на CPU) для дополнительной экономии.
• Модели 30B+ обязателен ZeRO-3. Часто комбинируется с Pipeline Parallelism или Tensor Parallelism (модельный параллелизм) для уменьшения объёма коммуникаций. Также полезен offload параметров на CPU/NVMe.

Важно при использовании ZeRO-3 на большом числе GPU (64+) коммуникационные накладные расходы могут свести на нет выигрыш в памяти. В таких случаях лучше использовать комбинацию ZeRO-3 с Pipeline Parallelism (DeepSpeed Pipe) или Tensor Parallelism (Megatron-LM).

8. Коммуникационные накладные расходы и компромиссы

• ZeRO-1 не добавляет новых коммуникаций, только стандартный AllReduce градиентов.
• ZeRO-2 заменяет AllReduce на ReduceScatter (такой же объём данных, но распределённый). Коммуникация практически не увеличивается.
• ZeRO-3 добавляет AllGather весов перед каждым слоем. Для модели размером M байт на N GPU каждая GPU передаёт M/N байт на слой. Если модель имеет L слоёв, общий объём переданных данных за шаг = L * (M/N) * (N-1) ≈ L*M. Это может быть значительно больше, чем в DP (где передаются только градиенты размером M). Поэтому ZeRO-3 чувствителен к латентности и пропускной способности сети.

Компромисс можно использовать ZeRO-3 с offload (параметры и optimizer states на CPU), что уменьшает объём коммуникаций, но добавляет задержки на transfer CPU↔GPU.

9. Связь с другими техниками распределённого обучения

• Data Parallelism (DP) — базовый подход, который ZeRO улучшает.
• Model Parallelism (MP) — разделение слоёв модели между GPU (Tensor Parallelism, Pipeline Parallelism). ZeRO можно комбинировать с MP для ещё больших моделей.
• Mixed Precision Training (FP16/BF16) — уменьшает размер весов и градиентов, что снижает нагрузку на память и коммуникации. ZeRO работает в паре с FP16.
• Gradient Accumulation — позволяет увеличить эффективный batch size без увеличения памяти. ZeRO не конфликтует с ней.
• DeepSpeed — библиотека, реализующая ZeRO, а также другие оптимизации (оптимизатор OneAdam, sparse attention и т.д.).

Пет-проект для закрепления

Задача Обучить модель GPT-2 (124M параметров) на задаче language modeling с использованием разных стадий ZeRO и сравнить использование памяти и скорость.

Инструменты PyTorch, DeepSpeed, Hugging Face Transformers, 2–4 GPU (можно использовать Google Colab с несколькими GPU или локально).

Шаги:

1. Установить DeepSpeed: pip install deepspeed.
2. Написать скрипт обучения на основе transformers.Trainer с интеграцией DeepSpeed.
3. Создать три конфигурационных файла DeepSpeed для стадий 1, 2, 3 (пример для stage 2):

   {
     "train_batch_size": 8,
     "gradient_accumulation_steps": 1,
     "fp16": {"enabled": true},
     "zero_optimization": {
       "stage": 2,
       "allgather_partitions": true,
       "allgather_bucket_size": 2e8,
       "overlap_comm": true,
       "reduce_scatter": true,
       "reduce_bucket_size": 2e8,
       "contiguous_gradients": true
     }
   }
   

4. Запустить обучение с каждым конфигом, замерить:
   • Пиковое использование памяти GPU (nvidia-smi или DeepSpeed report).
   • Время на один шаг (steps per second).
   • Возможность увеличить batch size.
5. Построить таблицу результатов.

Ожидаемый результат Вы увидите, что ZeRO-2 даёт значительную экономию памяти по сравнению с ZeRO-1, а ZeRO-3 — ещё больше, но может быть медленнее из-за AllGather. Для GPT-2 (124M) разница между ZeRO-2 и ZeRO-3 будет небольшой, но на более крупной модели (например, GPT-2 XL 1.5B) ZeRO-3 позволит обучаться на меньшем числе GPU.

Связь с другими вопросами

Навигация

• Предыдущий: 461
• Следующий: 463
• Индекс: 00. Индекс разборов