Disk Sparse Adam: обучаем гигантские графы в PyTorch без OOM

В статье на Habr ML представлена Open Source библиотека Disk Sparse Adam (DSA), решающая одну из самых болезненных проблем при обучении графных нейросетей – ошибку CUDA out of memory (OOM). Стандартный torch.optim.SparseAdam хранит параметры, градиенты и моменты полностью в RAM/VRAM, что при работе с эмбеддинг‑матрицами в миллионы узлов приводит к мгновенному заполнению памяти.

DSA меняет подход: вместо того, чтобы держать все массивы в видеокарте, она выносит состояния моментов на диск с помощью memory‑mapped файлов (mmap). При этом доступ к данным остаётся быстрым, а нагрузка на GPU существенно снижается. Библиотека полностью совместима с PyTorch, использует тот же API, что и SparseAdam, и требует лишь небольших изменений в конфигурации оптимизатора.

Ключевые особенности DSA:

• Дисковое хранение состояний (моменты, скоростные буферы) через mmap;
• Поддержка градиентных обновлений в онлайн‑режиме без полной загрузки данных в память;
• Минимальный оверхед – только несколько мегабайт RAM на процесс;
• Прозрачная интеграция с существующим кодом PyTorch, без необходимости переписывать модель.

В тестах авторов DSA позволила обучать графовые модели с более 10 млн узлов на видеокарте RTX 3060 (12 ГБ VRAM), где обычный SparseAdam уже на 2 млн узлах приводил к OOM. При этом скорость обучения упала менее чем на 5 % по сравнению с полностью в‑мемориным вариантом.

Библиотека доступна на GitHub под лицензией MIT, сопровождается подробной документацией и примерами. Для начала работы достаточно установить пакет disk-sparse-adam через pip и заменить в коде torch.optim.SparseAdam на dsa.DiskSparseAdam.

Ссылка на оригинальную статью: https://habr.com/ru/articles/1068202/