Как работает gradient checkpointing в DeepSpeed?

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

Gradient checkpointing (также известный как activation checkpointing) — это техника оптимизации памяти при обучении глубоких нейронных сетей, которая позволяет значительно снизить потребление видеопамяти (VRAM) за счёт пересчёта промежуточных активаций во время обратного прохода (backward pass). Вместо того чтобы хранить все активации всех слоёв, DeepSpeed разбивает модель на чанки и сохраняет активации только на границах этих чанков. При обратном распространении ошибки активации внутри каждого чанка пересчитываются заново. Это даёт выигрыш в памяти с O(L) до O(√L) или O(log L) в зависимости от стратегии, но добавляет около 30% вычислительных накладных расходов (overhead). Техника критически важна для обучения больших языковых моделей (LLM) размером 70B+ параметров, где без неё модель просто не помещается в память даже одного GPU.

1. Термин: Gradient Checkpointing (градиентный контрольный точки)

Gradient checkpointing — это метод уменьшения использования памяти при обучении нейронных сетей. Идея заключается в том, чтобы не хранить все промежуточные значения (активации), вычисленные во время прямого прохода (forward pass), а сохранять их только для некоторых «контрольных точек» (checkpoints). Во время обратного прохода (backward pass) активации между двумя контрольными точками пересчитываются заново из сохранённых значений.

Зачем это нужно При обучении глубоких моделей (например, GPT-3, LLaMA-70B) основным узким местом является не количество параметров, а память, необходимая для хранения активаций. Активации — это выходные значения каждого слоя после применения функции активации. Для модели с L слоями, размером батча B и скрытой размерностью H объём памяти для активаций составляет O(B * L * H). Для L = 96, B = 1, H = 12288 (как в GPT-3-175B) это ~1.4 ГБ на один слой, а всего ~135 ГБ — больше, чем помещается в один A100 (80 ГБ). Gradient checkpointing решает эту проблему.

Термин «Контрольная точка» (checkpoint) — это слой или группа слоёв, для которых активации сохраняются в памяти. Все остальные активации между контрольными точками не хранятся, а будут пересчитаны при необходимости.

2. Проблема: Память для активаций при обучении LLM

При обучении нейронной сети с использованием автоматического дифференцирования (autograd) требуется хранить все промежуточные активации для вычисления градиентов. Рассмотрим простую цепочку:

x → layer1 → a1 → layer2 → a2 → ... → layerL → loss

Прямой проход (forward): вычисляем a1, a2, ..., aL. Все они сохраняются в памяти.

Обратный проход (backward): используем сохранённые aL, aL-1, ... для вычисления градиентов по цепному правилу.

Потребление памяти

• Параметры модели: O(P) — обычно 2-4 байта на параметр (FP16/FP32).
• Активации: O(B * L * H) — может быть в 10-100 раз больше, чем параметры.
• Градиенты: O(P) — столько же, сколько параметры.
• Состояния оптимизатора (Adam): O(2P) — моменты первого и второго порядка.

Для LLM на 70B параметров в FP16:

• Параметры: 70B * 2 байта = 140 ГБ.
• Активации (B=1, L=80, H=8192): ~80 * 8192 * 2 байта ≈ 1.3 ГБ на слой, всего ~104 ГБ.
• Градиенты: ещё 140 ГБ.
• Оптимизатор: 280 ГБ.
• Итого: ~664 ГБ — невозможно разместить на одном GPU.

Gradient checkpointing позволяет сократить память для активаций с O(L) до O(√L) или O(log L), что даёт экономию в 5-10 раз.

3. Механизм работы gradient checkpointing в DeepSpeed

DeepSpeed реализует gradient checkpointing через разбиение модели на чанки (chunks). Количество чанков задаётся параметром num_checkpoints. Стратегия работы:

1. Разбиение модели Модель делится на num_checkpoints последовательных групп слоёв (чанков). Например, для модели из 96 слоёв и num_checkpoints=4 получим 4 чанка по 24 слоя.

2. Прямой проход (forward):

   • Вычисляются активации для первого чанка. Активации на выходе чанка (последнего слоя в чанке) сохраняются в памяти как контрольная точка.
   • Для второго чанка: берём сохранённую активацию с предыдущей контрольной точки, вычисляем все слои внутри чанка, сохраняем активацию на выходе чанка.
   • Повторяем для всех чанков.
   • Важно активации внутри чанка (промежуточные слои) не сохраняются.

3. Обратный проход (backward):

   • Начинаем с последнего чанка. У нас есть сохранённая активация на входе в чанк (контрольная точка) и градиент на выходе чанка.
   • Пересчёт (recompute): Запускаем прямой проход внутри чанка заново, используя сохранённую входную активацию. При этом вычисляются все промежуточные активации, которые теперь временно хранятся в памяти (только для этого чанка).
   • Вычисляем градиенты для слоёв чанка, используя пересчитанные активации.
   • Освобождаем пересчитанные активации после вычисления градиентов.
   • Переходим к предыдущему чанку, используя его сохранённую контрольную точку.

Память В любой момент времени хранятся только:

• Активации контрольных точек: num_checkpoints значений.
• Активации текущего пересчитываемого чанка: O(размер чанка).

Таким образом, память для активаций снижается с O(L) до O(num_checkpoints + L/num_checkpoints). Оптимум достигается при num_checkpoints = √L, давая O(2√L).

4. Реализация в DeepSpeed: конфигурация и API

DeepSpeed предоставляет простой способ включить gradient checkpointing через конфигурационный файл или программно.

Конфигурационный файл (ds_config.json):

{
  "zero_optimization": {
    "stage": 2
  },
  "gradient_checkpointing": {
    "enabled": true,
    "num_checkpoints": 4
  }
}

Программно (Python):

import deepspeed

model = MyLargeModel()
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params={
        "zero_optimization": {"stage": 2},
        "gradient_checkpointing": {"enabled": True, "num_checkpoints": 4}
    }
)

Параметр num_checkpoints

• По умолчанию: None (DeepSpeed выбирает автоматически, обычно √L).
• Рекомендуемые значения: от 2 до 8 для моделей с 30-100 слоями.
• Слишком мало чанков → мало экономии памяти.
• Слишком много чанков → большой overhead на пересчёт.

Важно DeepSpeed автоматически определяет, какие модули можно чекпоинтить. Обычно это TransformerLayer или Block. Пользователь может указать пользовательские модули через checkpointable_modules.

5. Компромиссы: память vs скорость

Gradient checkpointing — это trade-off между памятью и вычислительными затратами.

Overhead в 30% означает, что обучение занимает на 30% больше времени, но становится возможным на доступном оборудовании.

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

• Модель не помещается в память GPU → обязательно.
• Модель помещается, но с маленьким батчем → можно использовать для увеличения батча.
• Модель помещается с большим батчем → не нужно, так как overhead не оправдан.

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

DeepSpeed часто комбинирует gradient checkpointing с ZeRO для максимальной экономии.

7. Пример: расчёт экономии памяти для LLaMA-70B

Параметры LLaMA-70B:

• L = 80 слоёв
• H = 8192 (скрытая размерность)
• B = 1 (батч)
• FP16: 2 байта на значение

Без checkpointing

• Активации на слой: 8192 * 2 = 16 КБ (на одно значение). Но на самом деле активации — это тензоры размера (B, seq_len, H). Для seq_len=4096: 1 * 4096 * 8192 * 2 = 64 МБ на слой.
• Всего активаций: 80 * 64 МБ = 5.12 ГБ.
• Параметры: 70B * 2 = 140 ГБ.
• Градиенты: 140 ГБ.
• Оптимизатор: 280 ГБ.
• Итого: ~565 ГБ — нужно 8 A100 (80 ГБ).

С checkpointing (num_checkpoints=8):

• Чанки: 80 / 8 = 10 слоёв на чанк.
• Сохраняемые контрольные точки: 8 * 64 МБ = 512 МБ.
• Пересчёт: в каждый момент времени хранятся активации одного чанка (10 слоёв) = 640 МБ.
• Итого для активаций: ~1.15 ГБ (вместо 5.12 ГБ).
• Общая память: 140 + 140 + 280 + 1.15 = ~561 ГБ — всё ещё нужно 8 GPU, но теперь можно увеличить seq_len или батч.

Реальный сценарий На практике gradient checkpointing позволяет обучать LLaMA-70B на 4-8 A100 с батчем 1-2, тогда как без него потребовалось бы 16+ GPU.

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

Задача Обучить небольшую модель (например, GPT-2 с 12 слоями) на задаче генерации текста с и без gradient checkpointing, сравнить использование памяти и время шага.

Инструменты

• Python, PyTorch, Transformers (Hugging Face), DeepSpeed.
• GPU с 8-16 ГБ VRAM (например, T4 в Colab).
• nvidia-smi для мониторинга памяти.

Шаги:

1. Загрузите модель GPT-2 (124M параметров, 12 слоёв).
2. Создайте простой датасет (например, случайные тексты).
3. Настройте DeepSpeed с gradient_checkpointing: {enabled: false} и обучите 10 шагов, замеряя время и пиковую память.
4. Повторите с gradient_checkpointing: {enabled: true, num_checkpoints: 3}.
5. Сравните результаты.

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

• Без checkpointing: пиковая память ~2.5 ГБ, время шага ~0.5 сек.
• С checkpointing: пиковая память ~1.8 ГБ (экономия ~30%), время шага ~0.65 сек (overhead ~30%).
• Для модели с 24+ слоями экономия будет более заметной (50-70%).

Код для загрузки и обучения

import torch
import deepspeed
from transformers import GPT2LMHeadModel, GPT2Tokenizer

model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# Подготовка данных
texts = ["Hello, world!"] * 100
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=128)
dataset = torch.utils.data.TensorDataset(inputs["input_ids"], inputs["attention_mask"])

# Конфигурация DeepSpeed
ds_config = {
    "train_batch_size": 4,
    "gradient_checkpointing": {"enabled": True, "num_checkpoints": 3},
    "zero_optimization": {"stage": 2}
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params=ds_config
)

# Обучение
for epoch in range(1):
    for batch in torch.utils.data.DataLoader(dataset, batch_size=4):
        input_ids, attention_mask = batch
        outputs = model_engine(input_ids, attention_mask=attention_mask, labels=input_ids)
        loss = outputs.loss
        model_engine.backward(loss)
        model_engine.step()

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

10. Навигация

Навигация

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