Training

ZeRO

Zero Redundancy Optimizer — алгоритм DeepSpeed, устраняющий дублирование параметров между GPU.

Что такое ZeRO

ZeRO (Zero Redundancy Optimizer) — алгоритм распределённого обучения, разработанный в Microsoft Research (2019). В стандартном data parallelism каждый GPU хранит полную копию весов, градиентов и состояний оптимизатора (AdamW хранит момент и дисперсию — 2× объём весов). ZeRO устраняет эту избыточность, шардируя состояние обучения между GPU.

Три стадии ZeRO

ZeRO-1: шардируются только состояния оптимизатора (Adam momentum + variance). Каждый GPU обновляет только свою долю параметров. Экономия: ~4× VRAM при FP16.

ZeRO-2: шардируются состояния оптимизатора + градиенты. Каждый GPU накапливает только свои градиенты. Экономия: ~8×.

ZeRO-3: шардируются состояния оптимизатора, градиенты + параметры модели. Каждый GPU держит только 1/N весов; перед forward/backward нужен AllGather. Экономия: ~N×, но больше коммуникации.

Когда какую стадию выбирать

Стадия Применение
ZeRO-1 4× экономия без дополнительной коммуникации
ZeRO-2 Стандарт для большинства multi-GPU задач
ZeRO-3 Модели не помещаются на 1 GPU даже без оптимизатора

ZeRO с CPU Offload

ZeRO-Offload дополнительно переносит состояния оптимизатора (и иногда параметры) на CPU RAM. Это позволяет обучать 70B-модели на одном GPU с 80 ГБ VRAM при наличии 512 ГБ RAM, но обучение заметно замедляется (~10–50×).

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  }
}

Связанные термины

  • DeepSpeed — фреймворк, реализующий ZeRO
  • FSDP — аналог ZeRO-3 в PyTorch
  • gradient checkpointing — дополнительная техника экономии VRAM
  • VRAM — ресурс, высвобождаемый ZeRO

Готовы запустить GPU-задачу?

Запустить GPU-сервер