Вы запускаете финтюнинг Granite 4.0 H 1B на Tesla A100 с 40GB VRAM. Скрипт стартует, прогресс-бар ползет, и вдруг — бах! OOM. Out of memory. Все 40 гигабайт съедены, процесс убит, а вы сидите и думаете: "Как так? A100 же мощная карта". Знакомо? Это не баг, это фича PyTorch. И сегодня я покажу, как заставить эту "фичу" работать на вас.
Почему A100 пасует перед Granite 4.0
Granite 4.0 H 1B — это не просто большая модель. Это 1.1 миллиард параметров в формате bfloat16. Казалось бы, 2.2 гигабайта весов. Но в реальности финтюнинг требует в 15-20 раз больше памяти. Вот куда уходят ваши 40GB:
- Активации: на каждый токен в последовательности — свой тензор градиентов
- Оптимизатор states: AdamW хранит моменты для каждого параметра
- Градиенты: копия для каждого параметра модели
- Кэш CUDA: PyTorch жадничает и резервирует память "про запас"
Самое смешное (или грустное) — большая часть этой памяти не используется активно. PyTorch по умолчанию ведет себя как тот парень в отеле "все включено": берет все, что видит, на всякий случай. И не отдает, пока не придет горничная (OOM ошибка).
PYTORCH_CUDA_ALLOC_CONF: не магия, а кнопка сброса кэша
Переменная окружения PYTORCH_CUDA_ALLOC_CONF — это не серебряная пуля. Это скорее регулировочный винт в двигателе. Ее часто преподносят как "волшебную таблетку", но на самом деле это просто набор политик управления памятью CUDA.
Внимание: Неправильная настройка PYTORCH_CUDA_ALLOC_CONF может замедлить обучение в 2-3 раза. Это не бесплатный обед.
1 Разбираемся, что вообще можно настроить
Формат простой: PYTORCH_CUDA_ALLOC_CONF=ключ1:значение1,ключ2:значение2. Но дьявол в деталях. Вот основные ключи:
| Ключ | Что делает | Типичные значения |
|---|---|---|
| max_split_size_mb | Максимальный размер блока памяти для разделения | 32, 64, 128 |
| garbage_collection_threshold | Порог для запуска сборщика мусора | 0.8, 0.9 |
| expandable_segments | Разрешить расширение сегментов памяти | True, False |
| roundup_power2_divisions | Округление размеров выделений | 2, 4 |
Самая частая ошибка — ставить max_split_size_mb=2, потому что "так в интернете посоветовали". Для A100 с ее 40GB это самоубийство. Вы получите тысячи мелких блоков памяти, и аллокатор захлебнется в метаданных.
2 Конфигурация для Granite 4.0 на A100
После двух дней экспериментов и трех перезагрузок сервера я нашел рабочую конфигурацию:
Экспорт переменной: export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128,garbage_collection_threshold:0.9,expandable_segments:False
Почему именно так? Разберем по косточкам:
- max_split_size_mb:128 — A100 имеет пропускную способность памяти 1555 GB/s. Мелкие блоки убивают производительность. 128MB — оптимальный баланс между фрагментацией и эффективностью.
- garbage_collection_threshold:0.9 — запускаем сборку мусора, когда память заполнена на 90%. Раньше — слишком частые паузы. Позже — риск OOM.
- expandable_segments:False — запрещаем сегментам памяти расширяться. Звучит контр-интуитивно, но это предотвращает "захват" всей памяти одним большим тензором.
Проверяем, что настройка применилась:
python -c "import torch; print(torch.cuda.memory_summary())"
В выводе ищите строки "Allocator Settings" — там должны быть ваши параметры.
Unsloth: друг или враг?
Unsloth рекламирует ускорение финтюнинга в 30 раз. И да, он работает. Но у него свои аппетиты к памяти. По умолчанию Unsloth включает:
- 4-битные адаптеры (хорошо)
- Автоматическое смешение типов (хорошо)
- Градиентный чекпоинтинг (очень хорошо)
- Но также кэширование внимания в формате float32 (плохо для памяти)
Проблема в том, что для Granite 4.0 с его контекстом 8192 токенов кэш внимания в float32 занимает:
(8192 * 8192 * 4 байта) ≈ 268MB на один слой. Умножаем на 24 слоя — получаем 6.4GB только на кэш внимания. Это до того, как мы начали считать градиенты.
3 Заставляем Unsloth экономить память
Вот модифицированная инициализация модели через Unsloth:
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "ibm-granite/granite-4.0-1b-h",
max_seq_length = 2048, # Не 8192! Снижаем для финтюнинга
dtype = None, # Автовыбор (обычно bfloat16 для A100)
load_in_4bit = True,
token = "ваш_токен",
device_map = "auto",
attn_implementation = "flash_attention_2", # Обязательно!
use_gradient_checkpointing = True, # Сохраняет 40% памяти
_attn_implementation_internal = "sdpa", # Для PyTorch 2.3+
)
Ключевые моменты:
- max_seq_length = 2048 — да, вы теряете длинный контекст, но выживаете. Можно позже увеличить до 4096, когда оптимизируете память.
- attn_implementation = "flash_attention_2" — не опция, а необходимость. Снижает потребление памяти внимания в 3-5 раз.
- use_gradient_checkpointing = True — пересчитывает активации вместо их хранения. Тормозит на 30%, но экономит гигабайты.
Не используйте "load_in_8bit" с Unsloth для Granite 4.0. 8-битная загрузка конфликтует с 4-битными адаптерами Unsloth. Результат — тихий NaN в градиентах и сломанное обучение.
Трюки, о которых не пишут в документации
PYTORCH_CUDA_ALLOC_CONF и Unsloth — это только 60% решения. Остальные 40% — ручная настройка того, что разработчики считали "незыблемым".
Батч-сайз: не то, чем кажется
Вы ставите batch_size=2 и получаете OOM. Пробуете batch_size=1 — снова OOM. Что за чертовщина? А вот что: в Transformers есть gradient_accumulation_steps. И он умножает эффективный batch size.
При gradient_accumulation_steps=4 и batch_size=1 эффективный batch size = 4. Все градиенты за 4 шага накапливаются в памяти. Решение:
training_args = TrainingArguments(
per_device_train_batch_size = 1,
gradient_accumulation_steps = 1, # Сначала 1, потом увеличиваем
gradient_checkpointing = True,
optim = "paged_adamw_8bit", # Не "adamw_hf"!
fp16 = False, # Для A100 используйте bf16
bf16 = True,
...
)
PagedAdamW8bit — это AdamW с постраничной памятью. Он хранит состояния оптимизатора в CPU RAM и подгружает в GPU по мере необходимости. Медленнее на 5-10%, но экономит 4-6GB VRAM.
Очистка кэша между эпохами
PyTorch не очищает кэш автоматически. После первой эпохи в памяти остаются тензоры, которые "могут пригодиться". Они не пригодятся. Добавьте в цикл обучения:
if epoch % 1 == 0: # После каждой эпохи
torch.cuda.empty_cache()
torch.cuda.synchronize()
gc.collect()
Да, это тормозит обучение. Но лучше медленно, чем никогда (из-за OOM).
Диагностика: куда делась память?
Прежде чем что-то настраивать, нужно понять, что именно жрет память. Мой стек диагностики:
- nvidia-smi -l 1 — смотрим динамику использования памяти
- torch.cuda.memory_summary() — детальная статистика PyTorch
- fuser -v /dev/nvidia* — какие процессы держат GPU
- py-spy top --pid PID — профайлинг Python-кода
Частая находка: память утекает не в веса модели, а в промежуточные активации. Особенно в слоях внимания с большим контекстом.
Если вы работаете с похожими проблемами на другом железе, посмотрите мой гайд про MoE на T4 — там те же принципы, но другие ограничения.
Чего НЕ делать никогда
За годы борьбы с OOM я собрал коллекцию анти-паттернов. Вот топ-3:
| Ошибка | Почему плохо | Что делать вместо |
|---|---|---|
| PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:2 | Убивает производительность A100 на 70% | Используйте 64 или 128 для A100 |
| disable_cache() в Transformers | Ломает генерацию после обучения | Используйте use_cache=False в forward |
| Ручное управление памятью через torch.cuda.* | Конфликтует с Unsloth и gradient checkpointing | Доверьтесь фреймворкам |
Когда ничего не помогает
Бывает. Вы перепробовали все настройки, а OOM все равно выскакивает на 15-й минуте обучения. Что делать?
Первое — уменьшить max_seq_length до 1024. Да, это больно. Но 1024 токенов достаточно для большинства задач финтюнинга. Второе — использовать LoRA вместо полного финтюнинга. LoRA добавляет всего 0.1% параметров и экономит до 75% памяти.
Третье — рассмотреть альтернативные фреймворки. Если Unsloth не справляется, посмотрите в сторону DGX Spark или старый добрый Hugging Face PEFT.
И последнее — арендовать карту с большей памятью. Иногда проще заплатить за A100 80GB, чем две недели биться с оптимизациями. Кстати, о ценах: в статье "Где арендовать GPU дешевле DeepInfra" я сравнивал варианты.
Самое интересное: после всех этих оптимизаций вы можете обнаружить, что обучение идет даже быстрее, чем до OOM ошибок. Потому что аллокатор памяти перестал тратить 50% времени на поиск свободных блоков. Ирония судьбы — чтобы ускорить обучение, иногда нужно сначала его почти остановить.
Если столкнетесь с похожей проблемой на других моделях — например, пытаетесь запустить Granite 4 Small на ноутбуке — принципы те же. Масштабируете под доступные ресурсы.
И помните: OOM — это не ошибка, а предложение оптимизировать код. Иногда слишком настойчивое.
