Disk Sparse Adam: решение проблемы нехватки памяти при обучении огромных графов на потребительском железе
Обучение моделей на миллионах узлов традиционно требует сверхмощных серверов из-за высокого потребления видеопамяти оптимизаторами. Новая библиотека Disk Sparse Adam (DSA) решает эту проблему, перенося исторические состояния моментов (m и v) в файловую систему с использованием memory mapping. Это позволяет работать с гигантскими разреженными моделями на стандартных видеокартах, таких как NVIDIA RTX 3090 или 4090, почти не увеличивая нагрузку на GPU.
# Побеждаем OOM: как Disk Sparse Adam позволяет обучать миллионные графы на домашней видеокарте
Разработка нейросетевых моделей для обработки знаний (Knowledge Graphs) или рекомендаций сталкивается с фундаментальным ограничением архитектуры: памятью. Когда размер модели исчисляется десятками миллионов сущностей, стандартные инструменты глубокого обучения, такие как PyTorch, требуют колоссальных объемов видеопамяти (VRAM) для хранения не только весов, но и состояний оптимизатора, которые могут быть на порядок больше самих параметров.
Команда проекта Disk Sparse Adam разработала решение, позволяющее перенести эти состояния на хранилище, используя NVMe SSD в качестве виртуальной оперативной памяти. Это открывает возможности для проведения экспериментов на обычном ноутбуке или в бесплатных средах, таких как Google Colab.
Проблема: Когда VRAM становится узким местом
Рассмотрим типичный сценарий: обучение эмбеддингов для графа с 1,000,000 сущностей. Размер каждого вектора — 128 float32.
Объем данных: ``python # Примерный расчет num_entities = 1_000_000 embedding_dim = 128 weight_size_mb = num_entities * embedding_dim * 4 / 1024 # ~500 МБ
Сам по себе тензор весов занимает около 500 МБ. Однако оптимизаторы первого рода, такие как Adam, требуют хранения трех тензоров на каждый параметр: веса (w), первый момент (m) и второй момент (v).
Состояния оптимизатора могут занимать до 50-100% от объема самих весов. В нашем примере это еще около 500 МБ на GPU. Если же граф становится плотным или размерность векторов увеличивается до 1024, потребление VRAM быстро приближается к пределу современных видеокарт (обычно 12-24 ГБ), вызывая ошибку CUDA Out Of Memory (OOM).
Решение: Disk Sparse Adam (DSA)
Библиотека Disk Sparse Adam вводит подход Disk Sparse Optimization. Её ключевая идея заключается в использовании памяти страницы файла (mmap) для хранения исторических состояний оптимизатора.
Как это работает
1. Разделение данных: Мы храним текущие веса модели в GPU для быстрого доступа во время обучения. Однако состояния моментов (m и v) для разреженных параметров (например, из torch.nn.Embedding) выносятся на диск. 2. Memory Mapping: Библиотека создает файл на диске (оптимально — на NVMe SSD) и отображает его в адресное пространство процесса. Это позволяет обращаться к данным на диске с почти такой же скоростью, как к оперативной памяти, минимизируя задержки ввода-вывода. 3. Масштабируемость: Поскольку состояния хранятся на диске, их объем практически не ограничен размером видеопамяти. Вы можете обучать модель с миллиардами параметров, пока ваш диск способен их вместить.
Преимущества использования
* Доступность: Возможность запускать масштабные эксперименты на одном компьютере с одной видеокартой или в облачных сред типа Google Colab. * Экономия ресурсов: Не требуется аренда мощных серверов с множественными GPU для тестирования гипотез. * Производительность: При использовании современных NVMe дисков накладные расходы на чтение данных минимальны и не становятся значительным узким местом. * Специализация: Оптимизация заточена под разреженные градиенты, характерные для задач с большим количеством уникальных сущностей.
Бенчмарк: 1 млн сущностей на Kaggle Notebook
Для демонстрации эффективности было проведено тестирование в среде Kaggle Notebook, которая предоставляет ограниченную видеопамять. Задача состояла в обучении модели на 1 млн сущностей с размером вектора 128.
Условия эксперимента: * Количество сущностей: 1 000,000 * Размерность вектора: 128 * Общее состояние весов на диске: ~1.5 ГБ * Оптимизатор: DiskSparseRiemannianAdam
Результаты:
Запуск обучения показывал стабильную работу без переполнения памяти: `` 🚀 Старт обучения 1,000,000 сущностей на Kaggle GPU... Epoch 01/20 | Loss: 0.000100 | GPU VRAM Overhead: 0.00 MB Epoch 20/20 | Loss: 0.000028 | GPU VRAM Overhead: 0.00 MB
Ключевые метрики: * Пропускная способность: 134,212 образцов в секунду (Samples/Sec). Это очень высокая скорость для такой конфигурации. * Нагрузка на GPU: Прирост потребления видеопамяти составил 0.00 МБ сверх базового уровня. Все состояния оптимизатора успешно работали на диске. * Сходимость: Модель успешно сходилась к низкому значению функции потерь (Loss).
Этот результат подтверждает, что перенос состояний оптимизатора на диск позволяет обойти ограничения видеопамяти без существенной потери производительности.
Заключение
Проект Disk Sparse Adam представляет собой важное решение для исследователей и инженеров, работающих с большими данными. Используя открытый код под лицензией MIT, разработчики могут масштабировать свои модели, ограниченные только скоростью диска, а не видеопамятью.
Исходный код доступен на GitHub: [github.com/Assistentus/DSA](https://github.com/Assistentus/DSA). Разработчики приглашают сообщество тестировать библиотеку на различных задачах, таких как GNN, рекомендательные системы и задачи с большими графами.
Важно: Данная библиотека предназначена исключительно для разреженных (sparse) градиентов. Использование ее для плотных сверточных слоев или трансформеров не рекомендуется и неэффективно.
*Проект был представлен на платформе Habr, где также доступен полный код и обсуждения от сообщества.