Что такое DeepSpeed ZeRO-Offload и когда он полезен?

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

DeepSpeed ZeRO-Offload — это техника распределённого обучения, которая выгружает optimizer state и gradients из памяти GPU в CPU RAM (или NVMe), оставляя на GPU только веса модели и вычисления forward/backward. Это позволяет обучать модели с десятками миллиардов параметров даже на одном GPU с ограниченной памятью (например, 70B на 1× A100 80GB). Цена — снижение скорости обучения в 2–5 раз из-за передачи данных между CPU и GPU.

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

Обучение большой языковой модели (LLM) требует огромного объёма видеопамяти. Основные компоненты, занимающие память:

• Веса модели (model weights): N параметров × 2 байта (FP16) = 2N байт.
• Градиенты (gradients): ещё 2N байт.
• Optimizer state (состояние оптимизатора, например Adam): хранит два момента (m и v) — ещё 4N байт (в FP32) или 2N (в FP16 precision precision precision training|mixed precision).
• Активации (activations) для backward pass — могут быть огромными, но их можно сжимать через gradient checkpointing.

Итого для модели 70B параметров в mixed precision (FP16 веса + FP32 optimizer):

• Веса: 70B × 2 = 140 GB
• Градиенты: 70B × 2 = 140 GB
• Optimizer state: 70B × 4 × 2 = 560 GB (m и v в FP32)
• Всего: ~840 GB — невозможно разместить даже на 8× A100 80GB (640 GB).

Нужны методы распределения памяти между устройствами.

2. DeepSpeed ZeRO: три стадии оптимизации памяти

ZeRO (Zero Redundancy Optimizer) — библиотека от Microsoft, которая устраняет избыточность хранения данных между GPU. Выделяют три стадии:

Принцип: каждый GPU хранит только свою часть данных, а при необходимости получает недостающие части от других GPU через all-gather или reduce-scatter.

3. Что такое ZeRO-Offload?

ZeRO-Offload — расширение ZeRO, которое выгружает optimizer state и gradients не на другие GPU, а в CPU RAM (или NVMe SSD). GPU остаётся только:

• веса модели (частично, если ZeRO-3)
• вычисления forward/backward

CPU берёт на себя обновление весов (optimizer step) и хранение состояний.

Как это работает (упрощённо):

1. Forward/backward — на GPU, веса и активации там же.
2. После backward градиенты пересылаются в CPU (через PCIe).
3. На CPU выполняется optimizer step (Adam, SGD и т.д.) — обновляются веса.
4. Обновлённые веса возвращаются на GPU для следующей итерации.

Для ZeRO-3 с offload веса также могут частично храниться на CPU и подгружаться по мере необходимости.

Offload на NVMe

Если CPU RAM не хватает, можно использовать NVMe SSD как дополнительный уровень памяти (ZeRO-Infinity). Это ещё медленнее, но позволяет обучать модели до триллиона параметров.

4. Когда ZeRO-Offload полезен?

Основной сценарий — обучение модели, которая не помещается в видеопамять доступных GPU, но при этом есть достаточно CPU RAM.

Когда НЕ стоит использовать

• Если модель помещается в GPU без offload — offload только замедлит.
• Если CPU RAM недостаточна (например, 32GB для модели 70B не хватит).
• Если требуется максимальная скорость (offload добавляет latency).

5. Trade-offs: скорость vs память

Формула оценки времени шага

T_total = T_compute + T_comm + T_offload

где T_offload — время передачи данных CPU↔GPU (пропускная способность PCIe Gen4 ~32 GB/s, NVMe ~7 GB/s).

6. Сравнение с другими подходами

ZeRO-Offload часто комбинируют с gradient checkpointing и ZeRO-2/3 для максимальной экономии.

7. Пример конфигурации DeepSpeed для ZeRO-Offload

Файл ds_config.json:

{
  "train_batch_size": 8,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 1e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8
    }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 1e-5,
      "warmup_num_steps": 1000
    }
  }
}

Запуск:

deepspeed --num_gpus=1 train.py --deepspeed ds_config.json

8. Практические рекомендации

• Измеряйте bottleneck: используйте nvidia-smi и htop для отслеживания загрузки GPU/CPU.
• Настраивайте pin_memory: ускоряет передачу на CPU.
• Используйте overlap_comm: перекрывайте коммуникацию с вычислениями.
• Для очень больших моделей комбинируйте ZeRO-Offload с ZeRO-3 и NVMe offload (ZeRO-Infinity).
• Не забывайте про CPU RAM: для модели 70B нужно минимум 500GB CPU RAM (optimizer state 560GB + градиенты 140GB + веса 140GB ≈ 840GB, но часть может оставаться на GPU).

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

Задача Обучить модель LLaMA-7B (или её меньшую версию, например, TinyLLaMA-1B) на одном GPU с 8GB памяти, используя ZeRO-Offload. Сравнить скорость и потребление памяти с обычным обучением (без offload).

Инструменты

• Python, PyTorch, Transformers, DeepSpeed
• GPU с 8GB (например, RTX 3070/3080)
• Датасет: любой небольшой (например, Alpaca 52k)

Шаги:

1. Установить DeepSpeed: pip install deepspeed.
2. Написать скрипт обучения с использованием deepspeed.initialize.
3. Создать конфиг ZeRO-Offload (stage 2, offload optimizer на CPU).
4. Запустить обучение с deepspeed и замерить:
   • Потребление GPU памяти (nvidia-smi)
   • Время на шаг (через time или логи)
5. Запустить обучение без offload (если модель помещается — для TinyLLaMA-1B может поместиться) и сравнить.
6. Построить таблицу: память GPU, время шага, loss.

Ожидаемый результат

• С offload: модель обучается, GPU память ~6-7GB, время шага ~2-3 секунды.
• Без offload: либо OOM, либо ~1 секунда (если помещается).
• Вывод: offload позволяет обучать модели, которые иначе не влезают, ценой скорости.

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

Навигация

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