DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer) устраняет избыточность памяти в распределённом обучении, распределяя состояния оптимизатора, градиенты и параметры по GPU. ZeRO имеет три стадии, каждая из которых распределяет больший объём состояния, чем предыдущая. DeepSpeed также поддерживает выгрузку в память CPU или NVMe для дополнительной экономии. Каждая дополнительная стадия и уровень выгрузки снижают пиковое потребление памяти, но увеличивают объём меж-GPU коммуникации.

                  params        grads       opt states
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-1          │██████████│ │██████████│ │███░░░░░░░│  GPU 0
                │██████████│ │██████████│ │░░░███░░░░│  GPU 1
                │██████████│ │██████████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-2          │██████████│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │██████████│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │██████████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-3          │███░░░░░░░│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │░░░███░░░░│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │░░░░░░████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
  █ resident    ░ held on another GPU

ZeRO-2 разделяет градиенты и состояния оптимизатора, обеспечивая меньшие накладные расходы на коммуникацию по сравнению с ZeRO-3. Используйте ZeRO-3 только в том случае, если ваша модель не помещается на GPU с использованием ZeRO-2.

Установка

Установите DeepSpeed из PyPI или установите Transformers с дополнительным пакетом deepspeed.

pip install deepspeed
pip install transformers[deepspeed]

Если вы сталкиваетесь с ошибками установки, связанными с CUDA, обратитесь к документации DeepSpeed CUDA. Установка из исходного кода — более надёжный вариант, поскольку он соответствует вашему конкретному оборудованию и включает функции, ещё недоступные в выпуске на PyPI.

Настройка

Trainer интегрирует DeepSpeed через аргумент deepspeed, который принимает JSON-файл конфигурации. В качестве альтернативы можно использовать файл конфигурации Accelerate вместо TrainingArguments.

TrainingArguments

Файл конфигурации Accelerate

Используйте "auto" в конфигурации для значений, которые вы хотите, чтобы DeepSpeed заполнил из TrainingArguments. Если вы хотите явно указать значение, убедитесь, что используете одинаковое значение как для аргумента DeepSpeed, так и для TrainingArguments.

См. справочник DeepSpeed Configuration JSON для полного списка параметров конфигурации DeepSpeed.

{
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto" }
    },
    "fp16": { "enabled": "auto" }
}

Конфигурация DeepSpeed	TrainingArguments
`train_micro_batch_size_per_gpu`	`per_device_train_batch_size`
`gradient_accumulation_steps`	`gradient_accumulation_steps`
`optimizer.params.lr`	`learning_rate`
`fp16.enabled`	`fp16`

Передайте конфигурацию в аргумент deepspeed.

from transformers import TrainingArguments

args = TrainingArguments(
    deepspeed="path/to/deepspeed_config.json",
    ...
)

# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py

Этапы ZeRO

Выберите конфигурацию этапа ZeRO в качестве отправной точки.

ZeRO-1

ZeRO-2

ZeRO-3

Следующие поля важны для настройки обучения.

zero_optimization задаёт этап ZeRO.
```
{ "zero_optimization": { "stage": 3 } }
```
Установите аргументы размера батча и накопления градиентов в значение "auto". Если вы вручную зададите значения, не совпадающие с параметрами TrainingArguments, обучение продолжится с неверными значениями без каких-либо предупреждений.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 задаёт точность обучения. Установите значение "auto", чтобы оно соответствовало флагу bf16 в TrainingArguments.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save выполняет all-gather по всем GPU перед сохранением, восстанавливая полные тензоры из их фрагментов. Это аргумент ZeRO-3.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
Установите overlap_comm в true, чтобы скрыть задержку all-reduce за обратным проходом. Параметры allgather_bucket_size и reduce_bucket_size позволяют выбирать между скоростью коммуникации и объёмом памяти GPU. Меньшие значения экономят память, но замедляют коммуникацию.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer переносит оптимизатор в оперативную память CPU. Чтобы сэкономить ещё больше памяти, также перенесите параметры модели с помощью offload_param (только ZeRO-3). Установите pin_memory в true, чтобы ускорить передачу данных между CPU и GPU, но это закрепляет оперативную память, делая её недоступной для других процессов.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```
optimizer и scheduler по умолчанию используют оптимизатор и планировщик, настроенные в TrainingArguments. Установите значение "auto", чтобы DeepSpeed читал значения из TrainingArguments, если вам не нужен нативный для DeepSpeed оптимизатор, например LAMB.
```
{
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": { "total_num_steps": "auto", "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" }
    }
}
```
Если вы выгружаете оптимизатор, установите zero_force_ds_cpu_optimizer в false, чтобы использовать CPU-оптимизатор Adam от DeepSpeed.
```
{
    "zero_force_ds_cpu_optimizer": false
}
```

Чекпоинты

DeepSpeed сохраняет контрольные точки в разбитом формате, который нельзя загрузить напрямую с помощью from_pretrained(). Установите load_best_model_at_end() в True, чтобы Trainer отслеживал и перезагружал лучшую контрольную точку в конце обучения.

from transformers import TrainingArguments, Trainer

args = TrainingArguments(
    deepspeed="ds_config_zero3.json",
    load_best_model_at_end=True,
    ...
)
# after training, save a normal transformers checkpoint
trainer.save_model("./best-model")

Установка save_only_model=True пропускает сохранение полного состояния оптимизатора, что означает невозможность перезагрузки лучшей модели в конце обучения. Также установите stage3_gather_16bit_weights_on_model_save: true, чтобы восстановить полные веса из их фрагментов. Это необходимо для сохранения консолидированного артефакта модели 16 бит или state dict 16 бит с ZeRO-3. Transformers вызывает ошибку, если save_only_model=True используется вместе с load_best_model_at_end=True.

Для возобновления работы с различными конфигурациями параллелизма см. руководство DeepSpeed по Universal Checkpointing.

Дальнейшие шаги

Ознакомьтесь с главой Zero Redundancy Optimizer в The Ultra-Scale Playbook, чтобы узнать больше о том, как работает ZeRO.
Ознакомьтесь со статьями о ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, Democratizing Billion-Scale Model Training и Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub