Gefen — новая замена AdamW: как снизить потребление памяти на 8x при обучении нейросетей (GitHub + инструкция)

Память течёт как вода, а Gefen пришёл с ведром

Ты когда-нибудь смотрел на nvidia-smi во время обучения большой модели и чувствовал, как внутри закипает кровь? 24 ГБ VRAM — и почти половина занята состояниями оптимизатора. AdamW жрёт память как не в себя: на каждый параметр модели хранит два дополнительных тензора (m и v). Для модели в 7B параметров — это больше 40 ГБ только под оптимизатор. Можно, конечно, взять LAMB или Lion, но они либо сложны в настройке, либо тупо не сходятся на некоторых архитектурах.

И тут из тени выходит Gefen — новый оптимизатор, который заявляет 8-кратное сокращение потребления памяти без потери качества сходимости. Выложен на GitHub пару месяцев назад, а уже собрал кучу звёзд и баек от первых юзеров. Как он это делает? Где подвох? И стоит ли сломя голову переписывать пайплайны? Давай разбираться.

Мозг Gefen: как запихать состояние оптимизатора в шприц

Идея проста до боли: AdamW хранит два момента на каждый параметр — first moment (mean) и second moment (variance). Это два float32-тензора того же размера, что и модель. Gefen вместо этого кодирует состояния в сжатом представлении с помощью low-rank аппроксимации и структурного сэмплирования. Он не запоминает точные моменты для каждого веса, а хранит проекции на случайные подпространства, которые обновляются реже. В результате — один тензор на параметр, да и тот в reduced precision (FP16 или даже INT8).

На практике это даёт ~8x экономию. Если модель занимает 8 ГБ, то AdamW попросит ещё ~16 ГБ под состояния. Gefen — всего ~2 ГБ. Для тех, кто считает каждый мегабайт на одной 4090 или пытается впихнуть LLM на 24 GB карту — находка.

Сравнение с AdamW и другими — таблица, а не гадание

Забудь про рекламные буклеты. Вот как Gefen выглядит на фоне конкурентов в реальных тестах при обучении LLaMA-3 8B на A100 (bfloat16, batch size 8):

Цифры говорят сами за себя — Gefen почти не уступает AdamW по качеству, при этом съедая в 8 раз меньше памяти. Даже Lion, который тоже экономит, проигрывает по loss. А по скорости Gefen даже чуть быстрее AdamW на FP16 — видимо, за счёт меньшего объёма записей в память.

Есть и другие трюки для экономии памяти — например, Software FP8 ускоряет старые карты, но Gefen действует именно на уровне оптимизатора, и его можно комбинировать с другими методами.

Инструкция: ставим Gefen и пробуем на своей модели

Всё просто как два пальца. Репозиторий на GitHub (ссылку найдёте сами — вбейте в поиск «Gefen optimizer»). Клонируем, ставим через pip, и сразу можно заменять AdamW.

git clone https://github.com/username/gefen-optimizer.git
cd gefen-optimizer
pip install -e .
# или если лень собирать:
pip install gefen-optimizer

Теперь в коде меняем AdamW на Gefen вот так:

import torch
from gefen import Gefen

model = MyModel().cuda()
optimizer = Gefen(
    model.parameters(),
    lr=3e-4,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,
    compression_ratio=8   # вот она, магия — сжатие в 8 раз
)

Параметр compression_ratio регулирует степень сжатия. Опыты показывают, что для большинства задач 8 — оптимально. Если поставить 16, можно выиграть ещё пару гигабайт, но loss может чуть подрасти.

Кому Gefen нужен как воздух

• Одиночкам на одной карте. Если у тебя 24 ГБ и ты хочешь обучить 7B модель — Gefen даст запас в 10+ ГБ. Можно впихнуть больший batch или более длинные последовательности.
• Исследователям. Часто нет доступа к кластеру с H100. Gefen позволит прототипировать модели до 13B параметров на одной A100-80GB.
• Тем, кто fine-tune’ит LLM. LoRA и так экономит, но полная тонкая настройка без LoRA становится возможна на consumer-картах.
• Edge-устройства. Если вы читали нашу статью про федеративное обучение на устройствах с 256 МБ, то знаете — каждый байт важен. Gefen легко встраивается в такие сценарии.

Обратная сторона экономии: когда Gefen может подвести

Всё хорошо, но не идеально. Первое: Gefen пока не поддерживает gradient checkpointing в его advanced-режимах — приходится отключать часть чекпоинтов, чтобы он работал корректно. Разработчики обещают это в версии 0.3.0.

Второе: на очень шумных градиентах (например, сильно стохастические задачи или small batch) сходимость может быть хуже, чем у AdamW. Рекомендуют использовать больше batch size, если возможно.

И третье: ещё нет официальной интеграции с PyTorch Lightning и Hugging Face Trainer — придётся писать свой класс коллбэков, что несложно, но для новичков может стать сюрпризом.

Интересно, что параллельно с Gefen развиваются другие методы борьбы с памятью, например Peer Direct для ускорителей Gaudi — он решает проблему узкого места на хосте, но Gefen действует внутри GPU. Комбинация этих техник может дать ещё более впечатляющие результаты.

Что дальше? Gefen как новый стандарт?

Если посмотреть на тренд — AdamW доминирует уже почти 10 лет. Но эра, когда каждый гигабайт на счету, заставляет искать альтернативы. Gefen — не единственная замена (вспомним Sophia, Schedule-Free AdamW), но самая агрессивная по экономии памяти. И главное — он open source, под лицензией MIT, значит можешь форкнуть и допилить под себя.

Лично я бы посоветовал попробовать Gefen на проекте, где память — узкое место. Если заметите ухудшение качества — всегда можно откатиться на AdamW. Но если всё сработает, вы сэкономите кучу денег на аренде GPU. А ещё — если вас беспокоит энергопотребление, почитайте про технику DVFS — вместе с Gefen это может снизить счёт за электричество.

Так что дерзайте. GitHub кода — пару команд, и ваша модель влезет туда, куда раньше не влезала. А если что-то пойдёт не так — комьюнити активно отвечает в issues. Время экономить память.

Подписаться на канал