Почему training 70B модели требует optimizer sharding (ZeRO-3)?

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

Обучение модели с 70 миллиардами параметров (70B) в FP16 требует около 560 ГБ памяти только для хранения параметров, градиентов и состояний оптимизатора Adam. GPU|Один GPU (даже A100 80 ГБ) не вмещает такой объём. ZeRO-3 (Zero Redundancy Optimizer stage 3) решает проблему, шардируя (распределяя) все три компонента — параметры, градиенты и состояния оптимизатора — между GPU. Без sharding потребовалось бы минимум 8 GPU A100 80 ГБ для одной реплики модели, что неэффективно и дорого. ZeRO-3 позволяет обучать 70B модель на десятках или сотнях GPU с минимальным избыточным дублированием.

1. Термин: Memory footprint при обучении LLM

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

• Параметры модели (model parameters): веса нейронной сети.
• Градиенты (gradients): производные функции потерь по каждому параметру, нужны для шага оптимизации.
• Состояния оптимизатора (optimizer states): для Adam это моменты первого и второго порядка (mean и variance), а также иногда дополнительная копия параметров для точности.
• Активации (activations): промежуточные значения слоёв, сохраняемые для обратного распространения.

В этом разборе фокус на первых трёх, так как именно они шардируются в ZeRO-3. Активации обрабатываются отдельно (например, через gradient checkpointing).

2. Расчёт памяти для 70B модели в FP16

FP16 (16-bit floating point) — каждый параметр занимает 2 байта. 70 миллиардов параметров → 70e9 × 2 байта = 140 ГБ.

Adam optimizer хранит на каждый параметр два дополнительных значения: momentum (первый момент) и variance (второй момент). Оба обычно хранятся в FP32 (32-bit, 4 байта) для численной стабильности. Итого на каждый параметр: 2 (FP16 вес) + 4 (momentum) + 4 (variance) = 10 байт. Но часто Adam также хранит копию параметров в FP32 для точного обновления (precision precision precision precision mixed precision training). Тогда на каждый параметр: 2 (FP16 вес) + 4 (FP32 вес) + 4 (momentum) + 4 (variance) = 14 байт. В черновике указано упрощение: optimizer 280 ГБ (2 параметра на вес — momentum и variance, каждый FP32, итого 8 байт на вес, плюс копия градиентов? Разберёмся).

Реалистичный расчёт для mixed precision (FP16 forward/backward, FP32 master weights и optimizer states):

Однако на практике часто используют BF16 (bfloat16) вместо FP16, и master weights могут не храниться отдельно, если используется FP32 оптимизатор с FP16 градиентами. В черновике приведён более консервативный сценарий: модель 140 ГБ (FP16), optimizer 280 ГБ (только momentum и variance, без master weights), gradients 140 ГБ → 560 ГБ. Это соответствует случаю, когда master weights не хранятся, а обновление происходит в FP32 из FP16 градиентов (что менее стабильно). Для надёжности мы рассмотрим оба варианта, но суть та же: память намного превышает ёмкость одного GPU.

Вывод даже минимальная оценка 560 ГБ требует как минимум 8 GPU A100 80 ГБ (8 × 80 = 640 ГБ) для одной реплики модели, если не использовать шардирование. Это дорого и неэффективно — 8 GPU будут хранить полные копии всех данных, то есть 7/8 памяти тратится впустую.

3. Что такое ZeRO и его стадии

ZeRO (Zero Redundancy Optimizer) — техника распределённого обучения, разработанная Microsoft, которая устраняет избыточное дублирование данных между GPU. В стандартном Data Parallelism (DP) каждый GPU хранит полную копию модели, градиентов и оптимизатора. ZeRO разделяет эти данные между устройствами.

Три стадии ZeRO:

ZeRO-3 шардирует всё: параметры, градиенты и состояния оптимизатора. Каждый GPU хранит только свою часть (1/N от общего объёма). Во время forward/backward недостающие параметры собираются через all-gather коммуникацию. Градиенты после backward редуцируются через reduce-scatter.

4. Почему для 70B модели нужен именно ZeRO-3

Даже ZeRO-2 (шардирование градиентов и оптимизатора) оставляет полную копию параметров на каждом GPU. Для 70B модели это 140 ГБ (FP16) — больше, чем вмещает A100 80 ГБ. Значит, ZeRO-2 не подходит: параметры не помещаются на один GPU.

ZeRO-3 шардирует и параметры, поэтому каждый GPU хранит 140/N ГБ параметров. При N=64 это ~2.2 ГБ на GPU — легко. При N=8 это 17.5 ГБ — тоже влезает. Таким образом, ZeRO-3 — единственный способ уместить 70B модель на разумное количество GPU без использования model parallelism (тензорного или пайплайн-параллелизма).

Сравнение подходов для 70B модели (FP16, без учёта активаций):

На практике для 70B моделей часто комбинируют ZeRO-3 с Pipeline Parallelism (PP) и Tensor Parallelism (TP) (3D parallelism), но ZeRO-3 остаётся ключевым для оптимизатора.

5. Как работает ZeRO-3: коммуникация

При использовании ZeRO-3 каждый GPU хранит только свой шард параметров. Во время forward pass для каждого слоя необходимо собрать полные параметры этого слоя со всех GPU. Это делается через all-gather: каждый GPU отправляет свой шард всем остальным, и все получают полный набор. После использования слоя (например, в backward) шарды можно отбросить, чтобы освободить память.

Алгоритм шага обучения с ZeRO-3

1. Forward для каждого слоя:
   • All-gather параметров слоя (каждый GPU получает полные веса)
   • Вычислить forward через слой
   • Отбросить чужие шарды (оставить только свой)
2. Backward: аналогично, для каждого слоя:
   • All-gather параметров слоя (снова)
   • Вычислить backward, получить градиенты
   • Reduce-scatter градиентов (суммировать градиенты по всем GPU и оставить только свой шард)
3. Optimizer step каждый GPU обновляет только свой шард параметров, используя свой шард градиентов и состояний оптимизатора.

Коммуникационные затраты каждый forward/backward вызывает all-gather для каждого слоя. Это приводит к значительному объёму пересылок (порядка размера модели за шаг). Поэтому ZeRO-3 требует высокоскоростной сети (NVLink, InfiniBand) и оптимизаций (например, overlap communication with computation).

6. Альтернативы ZeRO-3

• FSDP (Fully Sharded Data Parallel) — реализация PyTorch, аналогичная ZeRO-3. Позволяет шардировать параметры, градиенты и оптимизатор. FSDP более гибко настраивает стратегию шардирования (например, можно не шардировать некоторые слои).
• DeepSpeed ZeRO — библиотека Microsoft, предоставляющая ZeRO-1/2/3, а также оптимизации для CPU offload (ZeRO-Offload) и NVMe offload (ZeRO-Infinity).
• Tensor Parallelism (TP) — разбивает матрицы весов на части, каждый GPU вычисляет часть слоя. Требует синхронизации на каждом шаге, но уменьшает память на параметры.
• Pipeline Parallelism (PP) — разбивает модель по слоям на разные GPU, каждый GPU обрабатывает свой набор слоёв. Уменьшает память, но вводит "баббл" (простои).

Для 70B моделей часто используют комбинацию: TP внутри узла (NVLink), PP между узлами, ZeRO-3 для оптимизатора.

7. Практические соображения

• Batch size: при использовании ZeRO-3 эффективный batch size = per-GPU batch size × число GPU. Для 70B моделей per-GPU batch size часто мал (1-2), поэтому общий batch size может быть недостаточным для сходимости. Используют gradient accumulation.
• Mixed precision: обязательно использование FP16/BF16 для экономии памяти. Adam states хранятся в FP32.
• CPU offload если GPU памяти всё равно не хватает, можно выгружать состояния оптимизатора на CPU (ZeRO-Offload). Это замедляет обучение, но позволяет обучать модели большего размера.
• Memory for activations ZeRO-3 не шардирует активации. Для 70B модели активации могут занимать сотни ГБ. Используют activation checkpointing (градиентные чекпоинты) — сохраняют только активации на границах слоёв, пересчитывая их при backward.

8. Пример кода: оценка памяти с помощью DeepSpeed

import deepspeed
import torch
from transformers import AutoModelForCausalLM

# Загружаем модель 70B (например, LLaMA-70B)
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-70b-hf")

# Конфигурация DeepSpeed ZeRO-3
ds_config = {
    "train_batch_size": 32,
    "gradient_accumulation_steps": 4,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 1e-5,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 0.01
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": True
        },
        "overlap_comm": True,
        "contiguous_gradients": True,
        "reduce_bucket_size": 5e8,
        "stage3_prefetch_bucket_size": 5e8,
        "stage3_param_persistence_threshold": 1e6,
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "gather_16bit_weights_on_model_save": True
    },
    "fp16": {
        "enabled": True,
        "auto_cast": True,
        "loss_scale": 0,
        "initial_scale_power": 16,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "gradient_clipping": 1.0,
    "steps_per_print": 100,
    "wall_clock_breakdown": False
}

# Инициализация DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params=ds_config
)

# Теперь model_engine — это обёртка с ZeRO-3
# Память распределена между GPU

Ожидаемый результат модель 70B успешно загружается на, скажем, 32 GPU A100 80 ГБ, каждый GPU использует ~40-50 ГБ памяти (с учётом активаций и оверхеда).

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

Задача Симулировать распределение памяти для 70B модели при разных стратегиях шардирования и визуализировать, сколько GPU потребуется.

Инструменты Python, библиотека psutil (для мониторинга памяти), torch.distributed (для имитации), matplotlib (для графиков).

Шаги:

1. Рассчитать теоретический memory footprint для 70B модели в FP16 (140 ГБ параметры, 140 ГБ градиенты, 280 ГБ optimizer states).
2. Написать функцию, которая для заданного числа GPU N и стадии ZeRO (0,1,2,3) вычисляет память на GPU.
3. Построить график зависимости памяти на GPU от N для каждой стадии.
4. Отметить лимит A100 80 ГБ и найти минимальное N для каждой стадии.
5. Добавить учёт активаций (например, 100 ГБ для batch size 1) и gradient checkpointing (уменьшает активации в 2-5 раз).

Ожидаемый результат Вы увидите, что ZeRO-3 позволяет обучать 70B модель на 16-32 GPU, тогда как ZeRO-2 требует >100 GPU (из-за параметров), а без ZeRO — невозможно. Код можно оформить как Jupyter notebook.

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

Навигация

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