ZeRO optimization — метод распределённого обучения, который делит параметры, градиенты и оптимизаторные состояния между устройствами, снижая потребление памяти.
Определение
ZeRO (Zero Redundancy Optimizer) — это метод масштабирования обучения больших моделей, разработанный в рамках DeepSpeed. Он устраняет избыточные копии параметров, градиентов и состояний оптимизатора, распределяя эти данные между устройствами. В отличие от data parallelism, где каждая копия модели хранит полный набор параметров, ZeRO делит данные так, чтобы каждый GPU держал только свою часть.
Эта стратегия позволяет обучать модели, которые значительно превышают объём памяти одного GPU, сохраняя при этом эффективность и точность обучения.
Как работает
ZeRO разбивает данные, связанные с обучением, на шардированные фрагменты. Метод имеет три основных уровня оптимизации:
- ZeRO-1 — шардирование оптимизаторных состояний (momentum, variance, экспоненциальные сглаживания).
- ZeRO-2 — шардирование градиентов между GPU.
- ZeRO-3 — полное шардирование параметров, градиентов и состояний оптимизатора.
На уровне ZeRO-3 каждая реплика модели фактически становится распределённой: ни один GPU больше не содержит полный набор параметров. При forward и backward pass необходимые шардированные фрагменты собираются через all-gather или reduce-scatter.
Ключевые элементы механизма:
- Шардирование — разделение параметров и состояний на части.
- On-demand сборка — GPU собирает только те части модели, которые нужны для текущей операции.
- Коммуникационные примитивы — использование all-gather, reduce-scatter и broadcast.
- Offloading — перенос части данных в CPU-память или NVMe для экономии GPU-памяти.
Где применяется
- Обучение моделей с десятками и сотнями миллиардов параметров.
- Тренинг LLM на кластерах с ограниченной GPU-памятью.
- Уменьшение потребления VRAM при фиксированном количестве GPU.
- DeepSpeed-пайплайны, комбинирующие ZeRO с tensor и pipeline parallelism.
- Обучение больших MoE-архитектур.
- Суперкомпьютерные тренинги с высокой коммуникационной пропускной способностью.
Практические примеры использования
ZeRO-3 используется при обучении крупных моделей, аналогичных GPT-3 и GPT-4-подобным архитектурам. Шардирование параметров позволяет обучать модели с объёмом параметров, который в 10–20 раз превышает доступную память одного GPU.
В проектах, ограниченных ресурсами, ZeRO помогает тренировать большие модели на меньшем количестве GPU, комбинируя шардирование с CPU-offloading для хранения частей оптимизатора.
При использовании DeepSpeed ZeRO работает в гибридных схемах: tensor parallelism делит матрицы, pipeline parallelism делит модель на стадии, а ZeRO устраняет дублирование параметров.
В MoE-системах ZeRO снижает потребление памяти роутерами и экспертов, распределяя их состояния по устройствам.
Ключевые свойства ZeRO optimization
- Отсутствие дублирования данных — параметры и состояния оптимизатора хранятся только один раз.
- Высокая масштабируемость — позволяет обучать модели, намного превышающие memory footprint одного GPU.
- Гибридность — работает совместно с другими типами параллелизма.
- Гибкое offloading — возможность переноса данных в CPU или NVMe.
Проблемы и ограничения
- Коммуникационная нагрузка — ZeRO требует частых коллективных операций.
- Усложнение инженерии — требуется точная настройка DeepSpeed-пайплайна.
- Увеличение latency — из-за on-demand загрузки параметров.
- Высокие требования к сети — эффективная работа требует NVLink или InfiniBand.
Преимущества и ограничения
- Плюс: резко уменьшает использование GPU-памяти.
- Плюс: позволяет обучать модели, недоступные при обычном data parallelism.
- Плюс: хорошо сочетается с параллелизмом по данным и модели.
- Плюс: поддерживает гибкие режимы offloading.
- Минус: коммуникационные накладные расходы.
- Минус: необходимость сложной настройки.
- Минус: сенситивность к архитектуре кластера.
- Минус: рост времени одного шага обучения.
Связанные термины
- Model parallelism
- Tensor parallelism
- Pipeline parallelism
- Data parallelism
- Distributed training
- DeepSpeed
- All-reduce
- Offloading