Сообщение torch.OutOfMemoryError при обучении с GRPO выглядит пугающе, но содержит точный диагноз: сколько памяти запрошено, сколько свободно и где именно произошёл сбой. В типичном случае — при попытке выделить 6,01 ГБ на карте с 2,72 ГБ свободной памяти — нехватка составляет около 3,3 ГБ. Трассировка стека указывает на конкретную операцию: forward pass при вычислении логитов для всего пакета токенов.
GRPO (Group Relative Policy Optimization) — метод обучения с подкреплением для языковых моделей, при котором модель генерирует несколько вариантов ответа на один промпт, а затем обновляет веса на основе сравнения их качества. Для быстрой генерации GRPO использует vLLM — отдельный движок вывода, который резервирует память GPU заранее и независимо от основного процесса обучения. Именно это делает управление памятью нетривиальным: два крупных потребителя работают одновременно.
| Параметр | Исходное значение | Оптимизированное значение | Эффект |
|---|---|---|---|
| MAX_SEQ_LENGTH | 1024 | 512 | Снижение активаций и KV-кэша |
| GPU_MEMORY_UTILIZATION | 0.6 | 0.5 | Экономия ~2,2 ГБ |
| PER_DEVICE_TRAIN_BATCH_SIZE | 4 | 2 | Вдвое меньше активаций |
| GRADIENT_ACCUMULATION_STEPS | 1 | 2 | Сохраняет эффективный batch=4 |
| NUM_GENERATIONS | 4 | 2 | Вдвое меньше памяти под генерации |
| LORA_RANK | 32 | 16 | Незначительное снижение памяти |
Память GPU при обучении с GRPO распределяется между тремя категориями. Первая — сама модель с адаптерами LoRA: для модели на 1 млрд параметров это обычно меньше 1 ГБ и редко становится причиной проблем. Вторая — резервирование vLLM: при GPU_MEMORY_UTILIZATION=0.6 на карте с 22 ГБ это 13,2 ГБ, которые исчезают ещё до начала обучения. Третья — активации при обучении, объём которых растёт пропорционально batch_size, длине последовательности, размерности скрытых представлений модели и числу слоёв. Для Gemma 3 1B (hidden_dim=2048, 18 слоёв) при batch=4 и seq=1024 один forward pass потребляет около 300 МБ — и это число умножается на NUM_GENERATIONS, то есть на количество вариантов, которые модель генерирует для каждого промпта.
Память под активации масштабируется как произведение batch_size × seq_length × hidden_dim × число_слоёв × 2 байта — и умножается на NUM_GENERATIONS.
Практический расчёт для конфигурации с MAX_SEQ_LENGTH=1024, GPU_MEMORY_UTILIZATION=0.6, batch_size=4 и NUM_GENERATIONS=4 на GPU с 22 ГБ даёт суммарное потребление 21–25 ГБ — очевидный перебор. Оптимизированная конфигурация для той же карты: MAX_SEQ_LENGTH=512, GPU_MEMORY_UTILIZATION=0.5, batch_size=2, NUM_GENERATIONS=2, LORA_RANK=16, LOAD_IN_4BIT=True. Итоговое потребление — около 17 ГБ, запас около 5 ГБ. При этом снижение batch_size с 4 до 2 компенсируется увеличением GRADIENT_ACCUMULATION_STEPS с 1 до 2: эффективный размер пакета остаётся равным 4, динамика обучения сохраняется.
Расставлять рычаги по приоритету стоит так: GPU_MEMORY_UTILIZATION даёт наибольший эффект на единицу изменения — снижение с 0.6 до 0.5 освобождает 2,2 ГБ. NUM_GENERATIONS умножает потребление памяти под активации, поэтому его уменьшение вдвое даёт линейный выигрыш. PER_DEVICE_TRAIN_BATCH_SIZE влияет на все активации сразу. MAX_SEQ_LENGTH сказывается и на активациях, и на KV-кэше vLLM. LORA_RANK вносит наименьший вклад и трогать его стоит в последнюю очередь.
Если после стандартных мер OOM сохраняется, есть дополнительные инструменты: агрессивное снижение GPU_MEMORY_UTILIZATION до 0.4, сокращение целевых модулей LoRA до минимального набора (q_proj и v_proj вместо всех проекций), а также переменная окружения PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, которая меняет стратегию аллокации памяти PyTorch. Мониторинг в реальном времени через nvidia-smi или torch.cuda.memory_allocated() позволяет видеть фактическое потребление на каждом шаге и точнее локализовать узкое место.