Нехватка CUDA-памяти при обучении с GRPO: как перестать гадать и начать считать

Недавно я собирал для заказчика модель обучения с подкреплением с использованием GRPO и Unsloth. Всё было настроено, набор данных был готов, и вижу:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.01 GiB. 
GPU 0 has a total capacity of 22.03 GiB of which 2.72 GiB is free.

Знакомо?

Я заметил вот что: когда большинство людей сталкиваются с ошибкой нехватки памяти (OOM), они начинают наугад менять параметры. Уменьшить размер пакета. Не помогло? Урезать длину последовательности вдвое. Всё ещё падает? Снизить ранг LoRA. Это метод проб и ошибок без реального понимания, почему что-то работает или не работает.

Я подхожу к этому иначе. Прежде чем что-либо менять, я хочу точно понять, куда уходит память. Тогда можно вносить точечные изменения, которые действительно решают проблему, не ухудшая без необходимости конфигурацию обучения.

Это руководство и есть такой подход, сведённый к практическому формату, который можно использовать уже сегодня.

Сообщение об ошибке

Это сообщение об ошибке – не просто шум, оно содержит всё необходимое. Давайте действительно его прочитаем:

Tried to allocate 6.01 GiB. 
GPU 0 has a total capacity of 22.03 GiB of which 2.72 GiB is free. 
Including non-PyTorch memory, this process has 19.29 GiB memory in use.

Вот что оно нам говорит:

Математика простая: нужно было 6,01 ГиБ, доступно было 2,72 ГиБ. Нам не хватает примерно 3,3 ГиБ.

Трассировка стека также показывает, где именно это произошло: в моём случае – во время выполнения get_per_token_logps_and_entropies при вычислении logits = model(**model_inputs).logits. Это прямой проход (forward pass), в котором считаются выходные логиты для всех токенов в пакете.

Теперь мы знаем, в чём проблема. Давайте разберёмся, что именно съедает память.

Куда на самом деле уходит память GPU в GRPO?

Прежде чем трогать какую-либо конфигурацию, нужно понять, кто потребляет память. При обучении с GRPO есть три основные категории:

1. Память модели: обычно небольшая

Для модели на 1 млрд параметров с LoRA общий объём обычно меньше 1 ГБ. Это не наша проблема.

2. Память vLLM для вывода: скрытый пожиратель ресурсов

GRPO использует vLLM для быстрой генерации. Вот что многие упускают: vLLM заранее резервирует фиксированную часть памяти GPU.

GPU_MEMORY_UTILIZATION = 0.6 # vLLM занимает 60% GPU

На GPU с 22 ГБ памяти это 13,2 ГБ, которые исчезают ещё до начала обучения. Часто это крупнейший потребитель памяти и при этом самый простой параметр для настройки.

3. Активации при обучении: главный виновник

Именно здесь обычно и возникают ошибки нехватки памяти. Память под активации масштабируется в зависимости от:

• размера пакета, PER_DEVICE_TRAIN_BATCH_SIZE;

• длины последовательности, MAX_SEQ_LENGTH;

• числа генераций, NUM_GENERATIONS;

• архитектуры модели: размерности скрытых представлений и числа слоёв.

Память под активации ≈ размер_пакета × длина_последовательности × размерность_скрытого_состояния × число_слоёв × 2 байта

Для Gemma 3 1B с hidden_dim=2048 и 18 слоями при batch=4 и seq=1024:

≈ 4 × 1024 × 2048 × 18 × 2 байта ≈ 300 МБ на один прямой проход

Но есть важный нюанс: GRPO генерирует NUM_GENERATIONS вариантов продолжения для каждого промпта. При NUM_GENERATIONS=4 вы умножаете это потребление памяти.

Процесс отладки: покажите расчёты

Давайте я подробно покажу, как именно диагностировал свою ошибку нехватки памяти.

Шаг 1. Перечислить всё

Моя исходная конфигурация:

MAX_SEQ_LENGTH = 1024
LORA_RANK = 32
GPU_MEMORY_UTILIZATION = 0.6
PER_DEVICE_TRAIN_BATCH_SIZE = 4
NUM_GENERATIONS = 4

Шаг 2. Рассчитать каждый компонент

У моего GPU 22 ГБ памяти. Я пытаюсь уместить в него 21–25 ГБ. Неудивительно, что всё упало.

Шаг 3. Найти самые сильные рычаги

Приоритет по степени влияния:

1. GPU_MEMORY_UTILIZATION – напрямую управляет тем, сколько памяти резервирует vLLM. Самый сильный одиночный рычаг.

2. NUM_GENERATIONS – умножает объём памяти, необходимый для сгенерированных продолжений.

3. PER_DEVICE_TRAIN_BATCH_SIZE – умножает объём памяти для всех активаций.

4. MAX_SEQ_LENGTH – влияет на активации и KV-кэш.

5. LORA_RANK – влияет слабее, но тоже вносит вклад.

Исправление: точечные изменения

На основе анализа вот моя оптимизированная конфигурация для GPU с 22 ГБ памяти:

# Конфигурация модели
MODEL_NAME = "google/gemma-3-1b-it"
MAX_SEQ_LENGTH = 512           # Уменьшено с 1024
LORA_RANK = 16                 # Уменьшено с 32
LOAD_IN_4BIT = True
GPU_MEMORY_UTILIZATION = 0.5   # Уменьшено с 0.6, экономит ~2,2 ГБ

# Конфигурация обучения
PER_DEVICE_TRAIN_BATCH_SIZE = 2  # Уменьшено с 4
GRADIENT_ACCUMULATION_STEPS = 2  # Увеличено, чтобы сохранить эффективный размер пакета
NUM_GENERATIONS = 2              # Уменьшено с 4

Новый расчет памяти

Запас памяти: 22 – 17 = ~5 ГБ свободно ✓

Сохранение динамики обучения

Обратите внимание: я не стал просто урезать всё подряд. Я увеличил GRADIENT_ACCUMULATION_STEPS:

Исходно: batch_size=4 × grad_accum=1 = эффективный размер пакета 4
Теперь: batch_size=2 × grad_accum=2 = эффективный размер пакета 4 ✓

Тот же эффективный размер пакета, похожая динамика обучения.

Короткая памятка: конфигурации под разные объёмы GPU-памяти

Вот что, по моему опыту, стабильно работает на разном оборудовании:

Это не магические числа, а стартовые значения, основанные на расчётах памяти выше. Настраивайте их под конкретную модель и набор данных.

Всё ещё получаете OOM? Примите экстренные меры

Если вы применили рекомендации выше, но всё ещё упираетесь в ограничения памяти:

1. Ещё сильнее уменьшите долю памяти для vLLM

GPU_MEMORY_UTILIZATION = 0.4  # Агрессивно, но работает

2. Сократите целевые модули LoRA

# Вместо того чтобы применять LoRA ко всем модулям, оставьте только самое необходимое
LORA_TARGET_MODULES = ["q_proj", "v_proj"]  # Уберите k_proj, o_proj и т. д.

3. Задайте конфигурацию памяти PyTorch

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

4. Следите за памятью в реальном времени

watch -n 1 nvidia-smi

Или в Python:

import torch
print(f"Allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB")
print(f"Reserved: {torch.cuda.memory_reserved()/1e9:.2f} GB")

Компромиссы

У каждого изменения есть цена. Нужно понимать, чем именно вы жертвуете:

Цель не в том, чтобы минимизировать потребление памяти, а в том, чтобы найти конфигурацию, которая даёт максимальное качество обучения в рамках ограничений вашего оборудования.

Главный вывод

Когда вы сталкиваетесь с OOM, перестаньте наугад подкручивать гиперпараметры. Вместо этого:

1. Прочитайте ошибку – она точно показывает, сколько памяти нужно и сколько есть в наличии.

2. Разложите потребителей памяти по категориям – резервирование vLLM, модель, активации.

3. Считайте перед изменениями – понимайте, куда уходит память.

4. Сначала беритесь за самые сильные рычаги – обычно это доля памяти для vLLM и размер пакета.

5. Сохраняйте то, что важно, – используйте накопление градиентов, чтобы сохранить эффективный размер пакета.

Разница между системной отладкой и случайной отладкой – это разница между решением проблемы за 10 минут и тремя часами раздражающих попыток.

Надеюсь, это сэкономит вам время при следующем запуске обучения с подкреплением.

Если после оптимизации памяти хочется глубже разобраться, как LLM устроены и как их встраивают в рабочие процессы, можно присмотреться к открытым урокам OTUS. Они бесплатные, проходят в рамках онлайн-курсов, а на занятиях можно задать вопросы преподавателям-практикам.

• 4 июня, 20:00. «Продвинутый анализ данных с помощью LLM». Записаться

• 15 июня, 20:00. «Интеграция ИИ-агентов в рабочую разработку: обвязка агента навыками и MCP». Записаться

Полный список бесплатных уроков по искусственному интеллекту, разработке и не только смотрите в календаре.