Побеждаем OOM: новая библиотека для обучения гигантских графов на обычной видеокарте

Все, кто обучал графные нейросети или работал с эмбеддингами на миллионы узлов, знают проблему: стандартный оптимизатор SparseAdam в PyTorch мгновенно переполняет оперативную и видеопамять, выбрасывая ошибку CUDA out of memory. Разработчики предлагают элегантное решение — открытую библиотеку Disk Sparse Adam (DSA), которая переносит состояния моментов на диск с помощью memory mapping. Это позволяет обходить ограничения даже на обычной домашней видеокарте.

Суть подхода в том, что DSA не держит все служебные данные в памяти, а размещает их в файле на диске и обращается к ним через mmap. Для разреженных моделей такая схема особенно эффективна: обновления касаются лишь малой доли параметров, поэтому отображаемые страницы памяти загружаются по мере необходимости. В итоге становится возможным обучать модели, которые раньше требовали серверных GPU с десятками гигабайт памяти.

Особенно актуально это для задач рекомендательных систем, анализа социальных сетей и работы с графами огромного масштаба. Судя по описанию, библиотека обещает сохранить скорость сходимости, свойственную оригинальному SparseAdam, но без критичного роста потребления ресурсов. Пока проект позиционируется как легковесный инструмент для исследователей и практиков, которым не хватает бюджетных мощностей.

Появление таких решений — важный шаг к демократизации машинного обучения. Если раньше пайплайны для гигантских графов были прерогативой крупных компаний, то теперь энтузиаст с обычным игровым компьютером может попробовать себя в амбициозных экспериментах. Остаётся надеяться, что библиотека быстро обрастёт документацией и примерами, а сообщество подтвердит её надёжность в реальных задачах. Пока же это одна из самых любопытных утилит для всех, кто сталкивался с жестоким лимитом памяти в PyTorch.