|
Professor Seleznov
|
Недавно я собирал для заказчика модель обучения с подкреплением с использованием 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.
Перевод
Ошибка PyTorch: не хватает памяти CUDA. Не удалось выделить 6,01 ГиБ. GPU 0 имеет общий объём памяти 22,03 ГиБ, из которых свободно только 2,72 ГиБ.
Знакомо?
Я заметил вот что: когда большинство людей сталкиваются с ошибкой нехватки памяти (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 ГиБ памяти.
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 млрд параметров с 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. Найти самые сильные рычаги
Приоритет по степени влияния:
- GPU_MEMORY_UTILIZATION – напрямую управляет тем, сколько памяти резервирует vLLM. Самый сильный одиночный рычаг.
- NUM_GENERATIONS – умножает объём памяти, необходимый для сгенерированных продолжений.
- PER_DEVICE_TRAIN_BATCH_SIZE – умножает объём памяти для всех активаций.
- MAX_SEQ_LENGTH – влияет на активации и KV-кэш.
- 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? Примите экстренные меры
Если вы применили рекомендации выше, но всё ещё упираетесь в ограничения памяти:
- Ещё сильнее уменьшите долю памяти для 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, перестаньте наугад подкручивать гиперпараметры. Вместо этого:
- Прочитайте ошибку – она точно показывает, сколько памяти нужно и сколько есть в наличии.
- Разложите потребителей памяти по категориям – резервирование vLLM, модель, активации.
- Считайте перед изменениями – понимайте, куда уходит память.
- Сначала беритесь за самые сильные рычаги – обычно это доля памяти для vLLM и размер пакета.
- Сохраняйте то, что важно, – используйте накопление градиентов, чтобы сохранить эффективный размер пакета.
Разница между системной отладкой и случайной отладкой – это разница между решением проблемы за 10 минут и тремя часами раздражающих попыток.
Надеюсь, это сэкономит вам время при следующем запуске обучения с подкреплением.
Если после оптимизации памяти хочется глубже разобраться, как LLM устроены и как их встраивают в рабочие процессы, можно присмотреться к открытым урокам OTUS. Они бесплатные, проходят в рамках онлайн-курсов, а на занятиях можно задать вопросы преподавателям-практикам.
- 4 июня, 20:00. «Продвинутый анализ данных с помощью LLM». Записаться
- 15 июня, 20:00. «Интеграция ИИ-агентов в рабочую разработку: обвязка агента навыками и MCP». Записаться
Полный список бесплатных уроков по искусственному интеллекту, разработке и не только смотрите в календаре.-Источник
|
