Как дебажить memory fragmentation в LLM сервере?

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

Memory fragmentation — одна из скрытых причин OOM и роста latency в LLM-серверах. Она возникает, когда менеджер памяти PyTorch разбивает зарезервированные блоки на фрагменты, и со временем даже при достаточном объёме свободной памяти невозможно выделить непрерывный блок для нового тензора. Дебаг включает мониторинг memory stats|torch.cuda.memory stats|memory_stats(), профилирование nsys и использование решений вроде Attention|Paged Attention (vLLM) или настройки CUDA-аллокатора. Разберём признаки, инструменты и пошаговую процедуру.

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

Memory fragmentation — состояние, при котором свободная память разбита на множество мелких несмежных участков, поэтому выделить большой непрерывный блок невозможно, хотя суммарный объём свободной памяти достаточен. В контексте GPU и LLM-серверов это происходит из-за разницы в жизненном цикле тензоров (активации, градиенты, KV-кэш разных запросов). PyTorch использует собственный аллокатор, который кэширует освобождённые блоки, чтобы избежать дорогих вызовов cudaMalloc. Но если блоки часто дробятся при частичном освобождении, память фрагментируется.

Фрагментация бывает двух типов:

• Внешняя — свободные блоки разбросаны, суммарно их много, но ни один не подходит по размеру.
• Внутренняя — блок выделен больше, чем нужно, и небольшая часть внутри него не используется (незначительна для GPU).

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

2. Признаки memory fragmentation

3. Инструменты для дебага

3.1 Базовый мониторинг

• nvidia-smi — общая used/free память на уровне карты. Не видит внутреннюю фрагментацию, но полезен для первичного обнаружения утечки или роста.
• nvidia-smi dmon — динамический мониторинг с метриками fb (frame buffer usage), bar1 (PCIe BAR1 memory).

3.2 PyTorch memory API

• torch.cuda.memory_summary() — подробный текстовый отчёт: показывает все сегменты, их размеры, количество свободных блоков, состояние кэша.
• torch.cuda.memory_stats() — словарь с численными значениями: allocated_bytes.all.current, reserved_bytes.all.current, active_bytes.all.current, а также количество сегментов и calls to cudaMalloc.

Пример:

import torch

# в цикле инференса
stats = torch.cuda.memory_stats()
print(f"Allocated: {stats['allocated_bytes.all.current'] / 1e9:.2f} GB")
print(f"Reserved:  {stats['reserved_bytes.all.current'] / 1e9:.2f} GB")
print(f"Active:    {stats['active_bytes.all.current'] / 1e9:.2f} GB")
print(f"Num OOM during run: {stats['num_alloc_retries']}")

Ключевое отношение — allocated / reserved: если меньше 0.5 при устойчивой нагрузке, фрагментация сильная.

3.3 Профилировщики

• nsys profile (NVIDIA Nsight Systems) — трассирует вызовы cudaMalloc/cudaFree. Позволяет увидеть, какие операции вызывают аллокации и как долго они длятся. Запуск:

  nsys profile -t cuda,osrt python server.py
  

  После профилирования открыть .nsys-rep в Nsight GUI, посмотреть раздел "CUDA Memory Allocations".
• py-spy — для профилирования CPU, косвенно помогает, если аллокации занимают много CPU.

3.4 Экспорт трассировки аллокатора

• PYTORCH_CUDA_ALLOC_CONF=record_shapes:1 — записывает формы тензоров для каждой аллокации.
• TORCH_DISTRIBUTED_DEBUG=INFO — логирование аллокаций при распределённой работе.

4. Разница между allocated, reserved, active

Фрагментация проявляется как большое количество свободных сегментов внутри reserved. memory_summary() показывает список сегментов: например, 20 сегментов по 10 MB и 1 сегмент 1 GB свободны — выделить 2 GB не получится.

5. Влияние фрагментации на производительность LLM-сервера

1. Увеличение latency — каждая аллокация в фрагментированном пуле может потребовать дефрагментации (слияния соседних блоков) или нового cudaMalloc, что занимает миллисекунды и блокирует GPU.
2. Снижение throughput — из-за того, что память расходуется неэффективно, приходится уменьшать batch size или длину контекста.
3. Неожиданные OOM — даже при среднем потреблении ниже лимита возникает ошибка CUDA out of memory. Это очень неприятный эффект для production.

Особенно остро проблема стоит для KV-кэша, память под который выделяется динамически для каждого запроса. Если после обработки запроса блоки кэша освобождаются не полностью (из-за внутренних стратегий vLLM или TGI), фрагментация нарастает.

6. Решения и профилактика

6.1 Использовать Paged Attention (vLLM)

Самое эффективное решение. Paged Attention разбивает KV-кэш на блоки фиксированного размера (pages) и управляет ими как операционная система виртуальной памятью. Блоки могут быть несмежными в физической памяти — фрагментация практически исчезает. vLLM практически не страдает от фрагментации.

6.2 Настроить аллокатор PyTorch

• PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync — использует асинхронный аллокатор (CUDA 11.4+). Он сам дефрагментирует память и позволяет выделять до 99% карты.

  export PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync
  python server.py
  

• expandable_segments:True (PyTorch 2.0+) — позволяет расширять существующие сегменты без создания новых фрагментов.

  export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  

• max_split_size_mb:128 — ограничивает максимальный размер блока, который может быть раздроблен. Меньшие значения снижают фрагментацию, но могут увеличить число аллокаций.

6.3 Периодическая очистка кэша

Вызов torch.cuda.empty_cache() освобождает все кэшированные блоки, возвращая их CUDA. Это дефрагментирует пул, но вызывает паузу (синхронизацию CUDA) и может снизить throughput на 1-5%. Использовать только при обнаружении растущей фрагментации (например, по триггеру allocated / reserved < 0.4):

if torch.cuda.is_available():
    before = torch.cuda.memory_allocated()
    torch.cuda.empty_cache()
    print(f"Freed {before - torch.cuda.memory_allocated()} bytes")

6.4 Рестарт сервера раз в сутки

Грубый, но рабочий способ для production, если другие решения внедрить сложно. Во время рестарта память полностью очищается. В комбинации с graceful shutdown (дообработка активных запросов) приемлемо.

6.5 Мониторинг и алертинг

Настроить метрики:

• allocated_bytes / reserved_bytes (если < 0.6 → предупреждение)
• num_alloc_retries (число повторных попыток аллокации — признак фрагментации)
• cuda_malloc_count — количество вызовов cudaMalloc в минуту.

По этим метрикам отправлять алерты (PagerDuty, Slack).

7. Пошаговая процедура дебага в production

1. Обнаружение проблемы — latency p99 начал расти, или появились OOM после нескольких часов работы. Проверяем nvidia-smi — used стабильно растёт.
2. Снимок состояния — подключаемся к серверу, выполняем torch.cuda.memory_summary(). Смотрим раздел "Allocated / Reserved".
3. Сравнение с baseline — запускаем на тестовом стенде с той же нагрузкой, но без фрагментации (например, сразу после рестарта). Сравниваем memory_stats.
4. Профилирование — запускаем nsys profile на 10 минутах нормальной нагрузки. Ищем частые вызовы cudaMalloc и их длительность.
5. Определение корня — если аллокации идут в основном от KV-кэша, переходим на vLLM. Если от промежуточных тензоров (активации, градиенты), настраиваем аллокатор.
6. Внедрение решения — меняем бэкенд или конфигурацию. Запускаем на канареечной доле трафика.
7. Валидация — проверяем метрики через сутки. allocated/reserved должно стать > 0.8, число вызовов cudaMalloc снизиться, latency p99 вернуться к норме.

8. Типичные ошибки при дебаге

• Ориентация только на nvidia-smi — не видит фрагментацию, только утечку.
• Путают allocated и reserved — смотрят на allocated (мал) и думают, что памяти достаточно, а на самом деле reserved почти исчерпан.
• Игнорирование влияния batch size — большой batch увеличивает размер тензоров, что ускоряет фрагментацию.
• Частый вызов empty_cache — убивает throughput и не решает проблему на корню.
• Не учитывают параллелизм — несколько потоков выделяют память одновременно; аллокатор не thread-safe в некоторых конфигурациях.

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

Задача: Создать скрипт, который симулирует фрагментацию при типичной инференс-нагрузке, регистрирует её и демонстрирует эффективность PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync.

Инструменты: Python, PyTorch, CUDA, matplotlib, psutil (для мониторинга host memory).

Шаги:

1. Создайте «эмулятор» LLM-сервера: он принимает случайные «запросы» (число токенов 128-2048) и для каждого выделяет KV-кэш (тензор размером [batch, layers, heads, seq_len, head_dim]), а через 10 шагов освобождает старые.
2. Вставьте вызов torch.cuda.memory_stats() каждые 50 шагов.
3. Запустите сначала с дефолтным аллокатором, запишите графики allocated/reserved и time per step.
4. Запустите повторно с os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'backend:cudaMallocAsync'.
5. Сравните графики и время выполнения.

Ожидаемый результат: Вы увидите, что при дефолтном аллокаторе allocated/reserved снижается с 0.8 до 0.3 за 1000 шагов, а время шага растёт. При cudaMallocAsync отношение остаётся > 0.8, а времена стабильны.

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

Навигация

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