Что такое activation offloading и когда он нужен?

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

Activation offloading — это техника перемещения промежуточных активаций (выходов слоёв нейросети) из GPU VRAM в CPU RAM (или NVMe-накопитель) во время forward pass и обратная загрузка в GPU во время backward pass. Она необходима для обучения моделей с очень длинным контекстом (>50k токенов), когда активации не помещаются в доступную видеопамять. Основная цена — замедление обучения в 3–5 раз из-за передачи данных по шине PCIe.

1. Термин: Активации (Activations)

Активации — это выходные тензоры каждого слоя нейронной сети после применения функции активации (например, ReLU, GeLU). Во время forward pass они вычисляются и сохраняются, потому что нужны для вычисления градиентов на этапе backward pass.

Почему активации занимают много памяти:

• Размер активаций пропорционален: batch_size × seq_len × hidden_dim × num_layers
• Для LLM (например, Llama 2 7B с hidden_dim=4096, 32 слоя) при seq_len=32k и batch_size=1 активации занимают ~32 × 4096 × 32 × 2 байта (fp16) ≈ 8.6 GB
• При seq_len=128k — уже ~34 GB, что превышает объём VRAM многих GPU (A100 80GB — предел)

Термин «VRAM» (Video Random Access Memory) — видеопамять GPU, используется для хранения параметров, градиентов, оптимизатора и активаций.

2. Проблема: Ограничение GPU памяти при длинном контексте

Современные LLM требуют огромных объёмов памяти при обучении на длинных последовательностях. Даже с gradient checkpointing (рекомпьютация активаций) может не хватать:

Gradient checkpointing — техника, при которой часть активаций не сохраняется, а пересчитывается заново во время backward. Это снижает память, но увеличивает время.

Однако при seq_len > 50k даже checkpointing может не спасти, потому что:

• Кэш KV (Key-Value) для attention растёт линейно с длиной контекста и занимает ~2 × batch_size × num_heads × seq_len × head_dim × num_layers байт
• Активации FFN (Feed-Forward Network) тоже велики

Тогда единственный выход — выгружать активации в CPU.

3. Что такое Activation Offloading

Activation offloading — это процесс:

1. Во время forward pass: после вычисления активаций каждого слоя (или группы слоёв) они асинхронно копируются из GPU в CPU RAM (или NVMe).
2. Во время backward pass: перед тем как слой понадобится для вычисления градиентов, активации prefetching'ом загружаются обратно в GPU.

Реализация обычно использует CUDA streams для перекрытия передачи данных и вычислений:

• Вычисления на одном stream, передача на другом
• Можно использовать pinned memory (закреплённая память CPU) для ускорения копирования

Термин «pinned memory» — страницы памяти CPU, которые не выгружаются на диск, что позволяет GPU обращаться к ним напрямую через DMA (Direct Memory Access).

4. Когда нужен Activation Offloading

Основные сценарии:

Когда не нужен

• Короткие контексты (<8k) — активации обычно помещаются
• Наличие GPU с большим VRAM (A100 80GB, H100 80GB, B200)
• Инференс без backward (активации не хранятся)

5. Цена и компромиссы

Основная цена — скорость передачи по PCIe:

• PCIe 4.0 x16: пропускная способность ~32 GB/s (теоретически), на практике ~25 GB/s
• PCIe 5.0 x16: ~63 GB/s
• NVMe SSD: ~3–7 GB/s (если offload на диск)

Для модели 7B с активациями ~10 GB на один шаг:

• Время копирования туда-обратно: 2 × 10 GB / 25 GB/s ≈ 0.8 сек
• Время одного шага без offloading: ~0.5–1 сек
• Итоговое замедление: 2–3x (если перекрывать вычисления и передачу)

На практике замедление 3–5x из-за:

• Наложения передачи на вычисления не всегда идеально
• Дополнительных накладных расходов на управление памятью
• Возможной фрагментации памяти

Компромисс: память vs скорость — offloading позволяет обучать модели, которые иначе не влезли бы, но ценой времени.

6. Альтернативы и сравнение

Важно activation offloading часто комбинируют с gradient checkpointing и FlashAttention для максимальной экономии памяти.

7. Реализация на практике

Пример псевдокода activation offloading в PyTorch с использованием CUDA streams:

import torch

class OffloadedLayer(torch.nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer
        self.offload_stream = torch.cuda.Stream()
        self.activations_cpu = []  # храним на CPU

    def forward(self, x):
        # Вычисляем активации
        act = self.layer(x)
        # Асинхронно копируем на CPU
        with torch.cuda.stream(self.offload_stream):
            act_cpu = act.to('cpu', non_blocking=True)
            self.activations_cpu.append(act_cpu)
        return act

    def backward(self, grad_output):
        # Загружаем активации обратно на GPU
        with torch.cuda.stream(self.offload_stream):
            act_gpu = self.activations_cpu.pop().to('cuda', non_blocking=True)
        # Синхронизируем потоки
        torch.cuda.current_stream().wait_stream(self.offload_stream)
        # Вычисляем градиенты
        return torch.autograd.Function.apply(self.layer, act_gpu, grad_output)

На практике используют готовые библиотеки:

• DeepSpeed — параметр activation_offloading в конфиге ZeRO-3
• FairScale — OffloadModel для offloading активаций
• PyTorch FSDP — cpu_offload для параметров, но не активаций

8. Связь с другими техниками оптимизации

Gradient accumulation — накопление градиентов за несколько шагов. Не снижает память активаций, но позволяет увеличить эффективный batch size.

Mixed precision training (fp16/bf16) — уменьшает размер активаций вдвое по сравнению с fp32. Обязательно использовать вместе с offloading.

FlashAttention — эффективная реализация attention, которая не хранит полную матрицу внимания (O(L²) → O(L)). Это снижает объём активаций attention, но не FFN.

Pipeline parallelism — распределение слоёв по разным GPU. Активации передаются между устройствами, что может быть альтернативой offloading.

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

1. Профилируйте память — используйте torch.cuda.memory_summary() и nvidia-smi для понимания, что именно занимает память.
2. Начинайте с gradient checkpointing — это проще и даёт хороший выигрыш.
3. Добавляйте activation offloading только если checkpointing не хватает.
4. Используйте pinned memory для CPU — ускоряет передачу.
5. Оптимизируйте размер активаций — уменьшайте hidden_dim, используйте LoRA, снижайте batch_size.
6. Комбинируйте с FlashAttention — это снижает нагрузку на attention.
7. Тестируйте на маленьких данных — проверьте, что offloading работает корректно.

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

Задача Реализовать activation offloading для обучения небольшой модели (например, GPT-2) на последовательностях длиной 8192 токенов и сравнить с baseline (gradient checkpointing).

Инструменты PyTorch, CUDA streams, GPT-2 из Hugging Face, torch.cuda.memory_summary().

Шаги:

1. Загрузите предобученный GPT-2 (124M параметров).
2. Создайте датасет из длинных текстов (например, книги из Project Gutenberg).
3. Обучите модель с gradient checkpointing (model.gradient_checkpointing_enable()) — замерьте время шага и пиковую память.
4. Реализуйте activation offloading для нескольких последних слоёв (где активации самые большие) с помощью CUDA streams.
5. Обучите с offloading — замерьте те же метрики.
6. Сравните результаты в таблице.

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

• Пиковая память GPU снизится на 20–40%
• Время шага увеличится в 2–3 раза
• Потери в качестве (perplexity) не будет, если offloading реализован корректно

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

Навигация

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