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

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

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

Ошибка PyTorch: не хватает памяти CUDA. Не удалось выделить 6,01 ГиБ. GPU 0 имеет общий объём памяти 22,03 ГиБ, из которых свободно только 2,72 ГиБ.

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

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

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

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

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

Не удалось выделить 6,01 ГиБ памяти.GPU 0 имеет общий объём памяти 22,03 ГиБ, из которых свободно только 2,72 ГиБ.С учётом памяти, занятой не только PyTorch, этот процесс уже использует 19,29 ГиБ памяти.

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

Математика простая: нужно было 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 с 22 ГБ памяти это13,2 ГБ, которые исчезают ещё до начала обучения. Часто это крупнейший потребитель памяти и при этом самый простой параметр для настройки.

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

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

  • размера пакета, PER_DEVICE_TRAIN_BATCH_SIZE;
  • длины последовательности, MAX_SEQ_LENGTH;
  • числа генераций, NUM_GENERATIONS;
  • архитектуры модели: размерности скрытых представлений и числа слоёв.

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

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

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

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

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

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

Шаг 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 ГБ памяти:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Или в Python:

Компромиссы

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

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

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

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

  1. Прочитайте ошибку– она точно показывает, сколько памяти нужно и сколько есть в наличии.
  2. Разложите потребителей памяти по категориям– резервирование vLLM, модель, активации.
  3. Считайте перед изменениями– понимайте, куда уходит память.
  4. Сначала беритесь за самые сильные рычаги– обычно это доля памяти для vLLM и размер пакета.
  5. Сохраняйте то, что важно, – используйте накопление градиентов, чтобы сохранить эффективный размер пакета.

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

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

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

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

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

Читать оригинал