Как вы управляете memory fragmentation при длительном раннинге LLM сервера?

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

Memory fragmentation (фрагментация памяти) — одна из ключевых проблем долгоживущих LLM-серверов, особенно при работе с KV cache (кэш ключей и значений) в трансформерах. Она приводит к неэффективному использованию GPU-памяти, снижению пропускной способности и, в конечном счёте, к OOM (out-of-memory) ошибкам. Основные методы борьбы: PagedAttention от vLLM, периодическая очистка кэша (torch.cuda.empty_cache()), настройка аллокатора через PYTORCH_CUDA_ALLOC_CONF, и проактивный рестарт сервера при превышении порога фрагментации (например, >30%).

1. Что такое memory fragmentation и почему она возникает при LLM инференсе

Memory fragmentation — это состояние, когда свободная память разбита на множество мелких несмежных блоков, из-за чего аллокатор не может выделить большой непрерывный участок, даже если суммарно свободной памяти достаточно.

При инференсе LLM основными потребителями памяти являются:

• Веса модели (загружаются один раз, статичны)
• KV cache (динамически растёт для каждого нового токена, сильно варьируется от длины последовательности)
• Промежуточные активации (при forward-проходе)

Фрагментация возникает из-за:

• Разнородных размеров запросов (разная длина входных и генерируемых последовательностей)
• Continuous batching (пакеты динамически формируются и распадаются)
• Освобождения памяти после завершения отдельных запросов (остаются «дыры»)

2. Типы фрагментации: внутренняя и внешняя

Для LLM-сервера критична внешняя фрагментация, так как KV cache требует больших непрерывных блоков (особенно для длинных последовательностей).

3. Как KV cache усугубляет фрагментацию

В стандартной реализации (например, Hugging Face Transformers) KV cache хранится как тензор размера [batch_size, num_heads, seq_len, head_dim] для каждого слоя. При генерации каждого нового токена этот тензор растёт по оси seq_len, что требует перевыделения памяти (reallocation) — старый тензор копируется в новый, больший блок, а старый освобождается. Это создаёт множество мелких освобождённых фрагментов.

Пример кода, иллюстрирующий проблему (упрощённо):

# Псевдокод: каждый шаг генерации перевыделяет KV cache
kv_cache = torch.zeros(1, 8, 0, 64)  # начальный пустой
for step in range(1000):
    new_cache = torch.zeros(1, 8, step+1, 64)  # новый больший тензор
    new_cache[:, :, :step, :] = kv_cache       # копирование
    kv_cache = new_cache                       # старый освобождается
    # -> множество освобождённых блоков разного размера

4. Решение vLLM: PagedAttention

PagedAttention — ключевая инновация vLLM, которая решает фрагментацию на уровне управления KV cache. Идея заимствована из виртуальной памяти ОС: KV cache разбивается на блоки фиксированного размера (pages), и каждый блок может быть размещён в любом свободном физическом фрейме. Логическая последовательность (запрос) использует таблицу страниц для отображения логических блоков в физические.

Как это работает

• Block size (размер блока) — обычно 16 или 32 токена.
• Page table (таблица страниц) — для каждого запроса хранит отображение логических номеров блоков на физические адреса.
• Allocation — при появлении нового токена выделяется новый блок (если текущий заполнен), но блок может быть взят из любого свободного физического фрейма.
• Deallocation — при завершении запроса освобождаются все его блоки, которые сразу становятся доступны для других запросов.

Преимущества

• Нет необходимости в непрерывных больших блоках — фрагментация практически устраняется.
• Память используется почти без потерь (только небольшая внутренняя фрагментация из-за неполностью заполненных блоков).
• Возможность разделять блоки между запросами (например, при prefix caching).

Сравнение

5. Дополнительные техники управления памятью

5.1 torch.cuda.empty_cache()

PyTorch использует кэширующий аллокатор (caching allocator). При освобождении тензора память не возвращается в ОС, а остаётся в пуле аллокатора для повторного использования. torch.cuda.empty_cache() принудительно очищает этот пул, возвращая неиспользуемые блоки в ОС.

Когда применять

• После завершения большого батча запросов.
• Периодически (раз в N запросов) — но не слишком часто, так как вызов дорогой.
• В комбинации с мониторингом: если фрагментация превышает порог.

Пример:

import torch
import gc

def cleanup_memory():
    gc.collect()  # сборка мусора Python
    torch.cuda.empty_cache()

5.2 PYTORCH_CUDA_ALLOC_CONF

Переменная окружения для тонкой настройки аллокатора PyTorch. Полезные параметры:

• max_split_size_mb:128 — максимальный размер освобождённого блока, который будет храниться в кэше (блоки больше этого размера сразу возвращаются в ОС). Уменьшение значения снижает фрагментацию, но увеличивает накладные расходы на выделение.
• expandable_segments:True — включает «расширяемые сегменты» (экспериментально в PyTorch 2.0+), которые позволяют динамически расширять выделенные блоки без перевыделения.

Пример использования

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:64,expandable_segments:True"
python serve_llm.py

5.3 Периодический рестарт сервера

Даже с PagedAttention со временем может накапливаться фрагментация из-за особенностей работы драйвера CUDA или утечек памяти в других компонентах (например, в Python-объектах). Рекомендуется:

• Мониторить memory fragmentation ratio (доля фрагментированной памяти).
• При превышении порога (например, >30%) выполнять graceful restart (перезагрузка сервера с сохранением состояния через очередь).
• Использовать health check с метрикой gpu_memory_fragmentation.

6. Мониторинг фрагментации

Для принятия решений нужно измерять фрагментацию. Основные инструменты:

6.1 nvidia-smi

Показывает общее использование памяти, но не фрагментацию.

6.2 torch.cuda.memory_summary()

Выводит детальную статистику аллокатора PyTorch: количество выделенных блоков, размер кэша, количество активных блоков, количество свободных блоков и их размеры.

Пример вывода

Allocated memory: 12.4 GiB
Cached memory: 14.2 GiB
Active memory: 11.8 GiB
Number of allocated blocks: 234
Number of cached blocks: 312
Fragmentation: 12.7% (estimated)

6.3 Кастомная метрика фрагментации

Можно вычислить как:

fragmentation = 1 - (max_free_block_size / total_free_memory)

Где max_free_block_size — размер самого большого непрерывного свободного блока, total_free_memory — общая свободная память (из torch.cuda.mem_get_info()).

Код для мониторинга

import torch

def get_fragmentation(device=0):
    free, total = torch.cuda.mem_get_info(device)
    # Получаем информацию о свободных блоках (требуется PyTorch >= 1.9)
    stats = torch.cuda.memory_stats(device)
    max_free_block = stats.get("max_split_size", 0)  # приблизительно
    if total == 0:
        return 0.0
    fragmentation = 1.0 - (max_free_block / free) if free > 0 else 0.0
    return fragmentation

7. Стратегии управления: комбинированный подход

Рекомендуемая стратегия для production LLM-сервера:

1. Использовать vLLM (или другой фреймворк с PagedAttention) как основу.
2. Настроить аллокатор PyTorch через PYTORCH_CUDA_ALLOC_CONF:
   • max_split_size_mb:64 (или 128 в зависимости от среднего размера KV cache).
   • expandable_segments:True (если поддерживается).
3. Периодически вызывать torch.cuda.empty_cache() после каждого N-го запроса (например, каждые 1000 запросов) или при превышении порога фрагментации >20%.
4. Мониторить фрагментацию в реальном времени и логировать метрику.
5. Настроить автоматический рестарт при фрагментации >30% (с предварительным drain запросов).

Пример конфигурации запуска vLLM

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:64,expandable_segments:True"
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-hf \
    --max-model-len 4096 \
    --gpu-memory-utilization 0.95 \
    --enable-prefix-caching

8. Сравнение подходов к управлению фрагментацией

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

• Начните с PagedAttention — это стандарт индустрии.
• Упомяните, что фрагментация — это не только GPU, но и CPU (например, при загрузке токенов), но для LLM сервера критична GPU.
• Расскажите про trade-off: уменьшение max_split_size_mb снижает фрагментацию, но увеличивает overhead на выделение памяти.
• Приведите пример из опыта: «В одном проекте мы заметили, что после 6 часов работы сервера latency выросла на 40% из-за фрагментации. Внедрение PagedAttention и периодической очистки кэша снизило latency до baseline и позволило работать неделями без рестарта».
• Не забудьте про мониторинг: без метрик вы не узнаете о проблеме.

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

Задача Написать симулятор memory fragmentation при LLM инференсе и протестировать стратегии управления.

Инструменты Python, PyTorch, CUDA (можно на CPU для простоты), matplotlib для визуализации.

Шаги:

1. Создайте класс Simulator, который эмулирует аллокатор с фрагментацией:
   • Выделяет блоки разного размера (имитация KV cache для запросов разной длины).
   • Случайным образом завершает запросы (освобождает блоки).
   • Отслеживает фрагментацию (по формуле выше).
2. Реализуйте три стратегии:
   • Baseline — ничего не делаем.
   • Periodic defrag — вызываем torch.cuda.empty_cache() каждые 100 шагов.
   • Paged-like — используем блоки фиксированного размера (аналог PagedAttention).
3. Запустите симуляцию на 1000 шагов, постройте график фрагментации во времени.
4. Сравните среднюю фрагментацию и количество OOM-ситуаций.

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

• Baseline покажет рост фрагментации до 40-50%.
• Periodic defrag — снижение до 20-30%, но с периодическими скачками.
• Paged-like — фрагментация <5% на протяжении всей симуляции.

Дополнительно Добавьте метрику «эффективное использование памяти» (доля реально используемой памяти от выделенной).

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

Навигация

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