DSA спасает память PyTorch, но переносит цену на SSD
Sh0ny
9 августа 20262 мин чтения
Коротко
Disk Sparse Adam позволяет обучать огромные sparse-эмбеддинги, не занимая гигабайты RAM и VRAM под состояния оптимизатора. Но это не ускорение: память экономится за счёт постоянного обмена данными с диском, а область применения ограничена разреженными градиентами.
У больших графовых моделей проблема часто не в самих параметрах, а в памяти оптимизатора. Для 10 млн сущностей с векторами размерностью 128 таблица весов занимает около 5,12 ГБ, а стандартный SparseAdam добавляет ещё примерно 10,24 ГБ на два состояния Adam.
Disk Sparse Adam (DSA) решает это довольно прямолинейно: хранит состояния m и v не в RAM или VRAM, а в файлах на диске, подключённых через mmap. Поскольку в каждом sparse-батче обновляется лишь небольшая часть сущностей, оптимизатор читает с диска только нужные индексы, обновляет их и записывает обратно.
Практический эффект понятен: модель, которая раньше не помещалась в память обычной рабочей станции, получает шанс запуститься на одной видеокарте или даже в бесплатном Colab. Это может быть полезно для графовых эмбеддингов, Knowledge Graph Embeddings, рекомендательных таблиц пользователей и товаров — везде, где параметры огромные, но на каждом шаге меняется лишь малая их часть.
Но важный trade-off здесь не в формуле Adam, а в месте хранения данных. DSA не делает обучение быстрее и не уменьшает общий объём состояний — он просто меняет дорогую память на дисковый ввод-вывод. В статье приведён синтетический запуск на 1 млн сущностей с заявленными 0 МБ дополнительной VRAM и 134 212 образцами в секунду, но сравнения со стандартным SparseAdam нет, поэтому эти цифры нельзя считать доказательством ускорения.
Ограничения существенные: желателен быстрый NVMe SSD, а HDD может стать узким местом. Оптимизатор рассчитан именно на sparse-параметры вроде Embedding и EmbeddingBag; для плотных свёрточных слоёв или трансформеров такой подход смысла не имеет. Кроме того, пример использования требует вручную читать веса по индексам и передавать градиенты в optimizer.step(), так что это скорее специальный инструмент для lookup-таблиц, чем универсальная drop-in замена.
Если перед вами стоит OOM из-за огромной embedding-таблицы, DSA выглядит разумным способом выиграть время и не арендовать сервер. Но если память помещается, сначала стоит измерить скорость I/O: готовы ли вы обменять нехватку VRAM на постоянную зависимость от SSD?