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

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

Memory fragmentation (фрагментация памяти GPU) — это ситуация, когда свободная видеопамять разбита на множество мелких фрагментов, из-за чего невозможно выделить непрерывный блок для нового тензора, хотя суммарно свободной памяти достаточно. При длительной работе LLM-сервера (дни/недели) фрагментация накапливается и может приводить к OOM-ошибкам даже при низком общем occupancy. Основные методы борьбы: использование PagedAttention (vLLM), настройка асинхронного аллокатора PyTorch, периодическая очистка кэша и мониторинг с автоматическим рестартом.

1. Термин: Memory Fragmentation (фрагментация памяти GPU)

Memory fragmentation — это дефект распределения памяти, при котором выделенные и освобождённые блоки разного размера создают «дырки» в адресном пространстве. В контексте GPU (CUDA) аллокатор по умолчанию (cudaMalloc) выделяет память большими пулами и раздаёт их под тензоры. После освобождения тензоров пулы не возвращаются в систему, а остаются в кэше аллокатора. Если тензоры были разного размера, пулы оказываются фрагментированными.

Последствия для LLM-сервера

• Невозможность выделить память под KV-cache для нового запроса.
• OOM (Out of Memory) при общем occupancy < 90%.
• Необходимость рестарта сервера.

Ключевые факторы, усугубляющие фрагментацию:

• Разнообразие длин последовательностей (разные batch size, разные размеры KV-cache).
• Долгое время работы без перезапуска.
• Использование динамических аллокаций (например, при search|beam search).

2. Почему фрагментация критична для длительного раннинга LLM сервера?

LLM-серверы (например, для чат-ботов или агентов) работают непрерывно. Каждый запрос выделяет память под KV-cache (ключи и значения attention), который хранится до завершения генерации. После ответа память освобождается, но не возвращается в систему — она остаётся в кэше аллокатора. Со временем пулы разного размера (из-за разной длины запросов) создают фрагментацию.

Пример:
Сервер обработал 10 000 запросов с длинами от 100 до 2000 токенов. После каждого запроса освобождался KV-cache разного размера. Через неделю аллокатор имеет 2 ГБ свободной памяти, но ни один непрерывный блок не превышает 100 МБ. Новый запрос с длиной 1500 токенов требует 300 МБ под KV-cache → OOM.

Агентный сценарий Агенты могут выполнять многошаговые рассуждения, удерживая контекст (state) между шагами. Это ещё дольше держит память занятой, увеличивая фрагментацию.

3. PagedAttention — решение от vLLM

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

Как это решает фрагментацию

• Выделение происходит страницами одинакового размера → нет фрагментации разного размера.
• Освобождённые страницы возвращаются в пул свободных страниц и могут быть использованы для любого другого запроса.
• Физическая память не фрагментируется, так как все блоки одного размера.

Дополнительные преимущества

• Возможность sharing страниц между запросами (например, при beam search).
• Эффективное использование памяти (до 4x больше запросов на одном GPU).

Недостатки

• Требует специальной реализации (vLLM, TensorRT-LLM).
• Накладные расходы на таблицу страниц (незначительны).

4. Альтернативы: PyTorch CUDA memory management

Если вы не используете vLLM (например, кастомный инференс на PyTorch), нужно управлять фрагментацией вручную. PyTorch использует собственный аллокатор CUDA (caching allocator), который по умолчанию пытается переиспользовать ранее выделенные блоки. Это хорошо для производительности, но плохо для фрагментации.

Основные инструменты

• torch.cuda.empty_cache() — освобождает все неиспользуемые кэшированные блоки, возвращая память в CUDA. Но это не решает фрагментацию пулов, которые уже выделены.
• torch.cuda.memory_summary() — показывает статистику использования памяти.
• torch.cuda.reset_peak_memory_stats() — сброс пиковых значений.

Проблема empty_cache() не дефрагментирует память. Он только возвращает неиспользуемые блоки, но пулы остаются фрагментированными.

5. Настройка PYTORCH_CUDA_ALLOC_CONF

PyTorch предоставляет переменную окружения PYTORCH_CUDA_ALLOC_CONF для настройки аллокатора. Для long-running серверов рекомендуется:

export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync

Что это даёт

• Использует асинхронный аллокатор CUDA (cudaMallocAsync), который лучше управляет фрагментацией.
• Позволяет освобождать память асинхронно, не блокируя вычисления.
• В некоторых версиях CUDA (11.4+) значительно снижает фрагментацию.

Другие параметры

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

Пример конфигурации для production

export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync,max_split_size_mb:128

6. Периодический restart сервера

Даже с PagedAttention или асинхронным аллокатором фрагментация может накапливаться (например, из-за утечек памяти в других компонентах). Простейший способ — graceful restart сервера при превышении порога фрагментации.

Практическая реализация

• Мониторинг фрагментации (см. раздел 7).
• Если фрагментация > 30% (или OOM-ошибки участились) → инициировать rolling restart.
• Для бесшовного рестарта использовать балансировщик (например, Nginx + несколько реплик).

Недостатки

• Потеря in-flight запросов (если не реализован drain).
• Время на перезагрузку модели (минуты).

Когда это оправдано

• Если нет возможности использовать vLLM.
• Для прототипов или low-load систем.

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

Чтобы вовремя заметить проблему, нужно мониторить состояние памяти GPU.

Инструменты

• nvidia-smi — показывает общее использование памяти, но не фрагментацию.
• torch.cuda.memory_stats() — детальная статистика PyTorch аллокатора.
• torch.cuda.memory_snapshot() — снимок всех выделенных блоков (можно визуализировать).

Ключевые метрики

• allocated_bytes.all.current — текущая занятая память.
• reserved_bytes.all.current — зарезервированная (кэшированная) память.
• active_bytes.all.current — память, используемая активными тензорами.
• Фрагментация = 1 - (максимальный свободный блок / общая свободная память). Вычисляется через memory_snapshot().

Пример кода для мониторинга

import torch

def get_fragmentation(device='cuda:0'):
    snapshot = torch.cuda.memory_snapshot(device)
    free_blocks = [b for b in snapshot if b['allocated_size'] == 0]
    if not free_blocks:
        return 0.0
    max_free = max(b['allocated_size'] for b in free_blocks)
    total_free = sum(b['allocated_size'] for b in free_blocks)
    return 1.0 - max_free / total_free if total_free > 0 else 0.0

# Вызов каждые N итераций
frag = get_fragmentation()
if frag > 0.3:
    print(f"Fragmentation {frag:.2%} > 30%, consider restart")

Интеграция с Prometheus/Grafana экспортируйте метрики через prometheus_client.

8. Сравнение подходов

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

1. Используйте vLLM или TensorRT-LLM — они решают проблему на уровне архитектуры.
2. Если vLLM недоступен, включите cudaMallocAsync и настройте max_split_size_mb.
3. Мониторьте фрагментацию каждые 5–10 минут.
4. Установите порог (например, 30%) для автоматического рестарта.
5. Для агентных систем (долгие сессии) используйте PagedAttention обязательно — без него KV-cache будет фрагментироваться за несколько часов.
6. Рассмотрите offloading части KV-cache на CPU при длинных контекстах (например, через FlexGen).

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

Задача Написать скрипт, который симулирует long-running LLM сервер на PyTorch и демонстрирует рост фрагментации. Затем применить методы борьбы.

Инструменты Python, PyTorch, CUDA, numpy, matplotlib.

Шаги:

1. Создайте простую модель (например, 2 слоя attention) и загрузите на GPU.
2. В цикле (1000 итераций) генерируйте случайные запросы разной длины (от 50 до 500 токенов). Для каждого запроса выделяйте KV-cache (тензор размера [batch, num_heads, seq_len, head_dim]), держите 10 итераций, затем освобождайте.
3. После каждых 100 итераций снимайте memory_snapshot() и вычисляйте фрагментацию.
4. Постройте график фрагментации от времени.
5. Повторите эксперимент с PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync и сравните.
6. Добавьте torch.cuda.empty_cache() каждые 50 итераций — оцените эффект.

Ожидаемый результат График, показывающий рост фрагментации без оптимизаций и её снижение при использовании асинхронного аллокатора и периодической очистки. Вывод: PagedAttention — лучшее решение, но в PyTorch можно смягчить проблему.

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

Навигация

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