Сколько VRAM нужно для локальной LLM: веса, KV cache, контекст и параллельные запросы

Введение​

Локальный запуск языковых моделей (LLM) требует тщательного планирования ресурсов, особенно объема видеопамяти (VRAM). Современные модели, такие как Llama-3.3-70B или Qwen-VL-72B, могут потреблять несколько гигабайт VRAM, что делает важным понимание того, как именно память используется в процессе инференса.

Это руководство подробно рассматривает, как рассчитать и оптимизировать использование VRAM при запуске LLM на GPU, с акцентом на ключевые компоненты: веса модели, KV cache, контекстную длину и параллельные запросы. Оно основано на данных, предоставленных vLLM, и включает практические рекомендации по минимизации потребления памяти и повышению эффективности.

Основные компоненты потребления VRAM​

Веса модели​

Веса модели — основной источник потребления VRAM. Размер весов зависит от числа параметров и формата хранения:

• FP32: 4 байта на параметр
  
• FP16: 2 байта на параметр
  
• INT8: 1 байт на параметр
  
• FP8: ~0.5 байта на параметр
  

Например, модель с 70 миллиардами параметров в FP16 будет занимать около 140 ГБ. Однако в реальности, особенно при использовании оптимизаций, такие значения могут быть значительно ниже. Например, использование FP8 или INT4 может сократить объем до ~35 ГБ.

KV cache​

KV cache (Key-Value cache) — механизм, используемый в трансформерах для хранения ключевых и значений на каждом шаге генерации. Это критически важно для аварегрессивной архитектуры, где каждое следующее слово зависит от предыдущих.

Каждый токен в KV cache занимает:

• Ключ: размер параметров × количество голов внимания × размер токена
  
• Значение: аналогично ключу
  

Формула расчета:

KV cache size = (batch_size * seq_len * num_layers * num_heads * head_dim) * 2 * dtype_size

Пример:

• batch_size = 1
  
• seq_len = 2048
  
• num_layers = 32
  
• num_heads = 32
  
• head_dim = 128
  
• dtype_size = 2 (FP16)
  

KV cache size ≈ 1 * 2048 * 32 * 32 * 128 * 2 = 262,144 KB ≈ 256 MB

Контекстная длина​

Контекстная длина — максимальное количество токенов, которое может быть обработано моделью. Чем длиннее контекст, тем больше требуется VRAM для KV cache.

В vLLM контекстная длина задается через max_model_len. При увеличении этой величины растет объем памяти, необходимый для хранения KV cache.

Параллельные запросы​

Параллельные запросы увеличивают нагрузку на VRAM. Каждый запрос требует выделения памяти под KV cache, и чем больше запросов в батче, тем больше памяти требуется.

vLLM поддерживает непрерывную батчинговую обработку, позволяя объединять запросы с разной длиной. Однако это может привести к переполнению памяти, если батч слишком большой.

Расчет необходимого VRAM​

Для расчета общего объема VRAM, необходимого для запуска модели, следует учитывать:

1. Веса модели
   
2. KV cache
   
3. Дополнительные метаданные
   
4. Кэширование и оптимизация
   

Пример расчета​

Рассмотрим модель Llama-3.1-8B-Instruct с параметрами:

• Размер весов: 8 миллиардов параметров
  
• Формат хранения: FP16
  
• Размер весов: 8B * 2 байта = 16 ГБ
  
• KV cache: 1 * 2048 * 32 * 32 * 128 * 2 = 256 MB
  
• Дополнительные метаданные: ~100 MB
  

Итого: ~16.3 ГБ

Если модель работает с батчом из 4 запросов, каждый длиной 2048 токенов:

• Общий KV cache: 4 * 256 MB = 1 ГБ
  
• Общий объем: ~17.3 ГБ
  

Ограничения и рекомендации​

• Минимальная память: Для большинства моделей рекомендуется иметь не менее 8 ГБ VRAM.
  
• Максимальная память: Для моделей с 70B параметров рекомендуется 24 ГБ и выше.
  
• Оптимизация: Использование квантования (INT4, FP8) может сократить потребление до 1/4.
  

Оптимизации памяти в vLLM​

vLLM предоставляет ряд механизмов для оптимизации использования VRAM:

Уровни оптимизации (-O)​

vLLM предоставляет четыре уровня оптимизации:

• -O0: Без оптимизации. Быстрый запуск, но низкая производительность.
  
• -O1: Быстрая оптимизация. Простая компиляция и быстрые фьюжны.
  
• -O2: Стандартная оптимизация. Дополнительные фьюжны и полные CUDA графы.
  
• -O3: Агрессивная оптимизация. Может включать экспериментальные оптимизации.
  

Reuse compile cache​

vLLM сохраняет результаты компиляции в ~/.cache/vllm. Это позволяет избежать повторной компиляции при повторном запуске.

Skip memory profiling​

Использование флага --kv-cache-memory позволяет пропустить профилирование памяти. Однако это может привести к снижению производительности или OOM при неправильной оценке.

Serve without CUDA graphs​

Флаг --enforce-eager отключает CUDA графы, что ускоряет запуск, но снижает стабильную производительность.

Управление KV cache​

Preemption​

Если память недостаточна для обработки всех запросов, vLLM может предварительно завершить некоторые запросы (preempt) и пересчитать их позже.

WARNING 05-09 00:49:33 scheduler.py:1057 Sequence group 0 is preempted by PreemptionMode.RECOMPUTE mode because there is not enough KV cache space.

Решения при preemption​

1. Увеличить gpu_memory_utilization
   
2. Уменьшить max_num_seqs или max_num_batched_tokens
   
3. Увеличить tensor_parallel_size
   
4. Увеличить pipeline_parallel_size
   

Chunked prefill​

Chunked prefill — механизм, позволяющий обрабатывать большие запросы частями. Это улучшает баланс между вычислениями (prefill) и памятью (decode).

Настройка​

max_num_batched_tokens = 16384

Эффекты​

• Лучшая latency: меньшие значения (2048) уменьшают задержку между токенами.
  
• Лучшая TTFT: большие значения (8192) ускоряют первый токен.
  

Параллелизм в vLLM​

Tensor parallelism​

Распределение весов модели между несколькими GPU. Полезно для моделей, которые не помещаются на один GPU.

tensor_parallel_size = 4

Pipeline parallelism​

Распределение слоев модели между GPU. Позволяет уменьшить нагрузку на каждый GPU.

pipeline_parallel_size = 2

Data parallelism​

Репликация модели на несколько GPU. Полезно для масштабирования пропускной способности.

data_parallel_size = 2

Expert parallelism​

Специализированный параллелизм для MoE моделей. Распределяет эксперты между GPU.

enable_expert_parallel = True

NUMA binding​

На многопроцессорных серверах NUMA binding помогает улучшить производительность, связывая процессы с ближайшими NUMA узлами.

--numa-bind

Batch-level DP​

Для многомодальных моделей (например, Qwen-VL) можно использовать batch-level DP, чтобы распределить входные данные по GPU вместо весов.

mm_encoder_tp_mode = "data"

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

1. Используйте FP8 или INT4, если возможно, чтобы сократить объем памяти.
   
2. Настройте max_num_batched_tokens в зависимости от ваших целей: latency или TTFT.
   
3. Мониторьте preemption через Prometheus.
   
4. Используйте tensor_parallel_size для распределения модели на несколько GPU.
   
5. Проверяйте gpu_memory_utilization при высокой нагрузке.
   
6. Используйте --numa-bind на многопроцессорных серверах.
   

Заключение​

Правильное планирование VRAM для локального запуска LLM требует понимания различных факторов: размера модели, KV cache, контекста и параллельных запросов. Использование оптимизаций vLLM, таких как chunked prefill, tensor parallelism и NUMA binding, позволяет значительно повысить эффективность использования памяти и производительность.

Понимание этих механизмов позволяет не только избежать OOM, но и оптимизировать производительность в зависимости от задачи. Для максимальной эффективности рекомендуется использовать комбинированные подходы, такие как квантование + параллелизм + оптимизация батчинга.

Источники​

• Optimization and Tuning - vLLM
  
• CUDA Toolkit Documentation
  
• vLLM