PyTorch: оптимизация памяти через диск сократила расход GPU с 10 ГБ до мегабайт

Проблема знакома каждому, кто работает с большими графами знаний или рекомендательными системами: PyTorch падает с ошибкой «CUDA out of memory», хотя сама модель в память помещается. Виноват не размер модели, а оптимизатор, который хранит два «момента» (историю градиентов для ускорения обучения) для каждого параметра. Автор открытого пакета DSA, распространяемого под лицензией MIT, предложил выносить эти моменты на диск через механизм mmap (отображение файлов в память, когда операционная система подгружает данные страницами по запросу). Для исследователей и стартапов в России, где аренда мощных GPU-серверов обходится дорого, а бюджетные видеокарты вроде RTX 3090 или 4090 есть у многих, это практичный способ запускать эксперименты без облачной инфраструктуры.

Почему стандартный SparseAdam съедает всю память?

Возьмём конкретный расчёт. Модель для графа знаний: 10 миллионов сущностей, размерность вектора 128, формат float32.

• Сами параметры (веса): 10 000 000 x 128 x 4 байта = 5,12 ГБ
• Первый момент (m): ещё 5,12 ГБ
• Второй момент (v): ещё 5,12 ГБ

Итого оптимизатор отнимает 10,24 ГБ сверх самой модели. Общий расход: 15,36 ГБ. Видеокарта с 16 ГБ видеопамяти уже не тянет, а если увеличить размерность до 256 или взять граф на 50 миллионов узлов, PyTorch падает с ошибкой CUDA out of memory, либо система зависает, заполнив весь файл подкачки.

При этом в каждом мини-батче (порции данных для одного шага обучения) обновляются не все 10 миллионов узлов, а, например, 10 000. Держать в дорогой памяти GPU историю для остальных 9 990 000 узлов бессмысленно.

Что понадобится

• Видеокарта: любая с поддержкой CUDA, подойдёт даже бесплатный Google Colab
• Диск: NVMe SSD (на обычном HDD будет медленно)
• Python и PyTorch: стандартная установка
• Пакет DSA: устанавливается из исходного кода на GitHub (лицензия MIT, открытый код)
• Время: замена оптимизатора в существующем коде занимает несколько минут

Пошаговая инструкция

1. Установите пакет DSA из репозитория на GitHub. Пакет работает как прямая замена (drop-in replacement) стандартного оптимизатора: переписывать архитектуру модели или загрузчик данных не нужно.

2. Создайте эмбеддинги (векторные представления узлов графа) с параметром sparse=True:

import torch
embedding = torch.nn.Embedding(10_000_000, 128, sparse=True)

1. Инициализируйте оптимизатор DSA вместо стандартного SparseAdam, указав папку на диске для хранения состояний:

from dsa.optimizer import DiskSparseRiemannianAdam

optimizer = DiskSparseRiemannianAdam(
    params={"emb": embedding.weight},
    lr=0.001,
    k=0.0,       # 0.0 для евклидова пространства, 1.0 для гиперболического
    disk_dir="./opt_cache"
)

Параметр k управляет геометрией пространства: при k=0.0 оптимизатор работает в обычном евклидовом пространстве, при k=1.0 используется шар Пуанкаре (гиперболическое пространство, которое лучше подходит для иерархических структур вроде деревьев категорий).

1. Запустите цикл обучения, передавая индексы батча и градиенты:

for batch_indices in dataloader:
    idx_np = batch_indices.numpy()
    weights_np = optimizer.state_files["emb"]["w"][idx_np].copy()
    current_weights = torch.from_numpy(weights_np).requires_grad_(True)

    loss = compute_loss(current_weights)
    loss.backward()

    optimizer.step(updates={"emb": (batch_indices, current_weights.grad)})

optimizer.shutdown()

1. Вызовите optimizer.shutdown() после завершения обучения, чтобы корректно завершить фоновый поток записи на диск.

Что делать с этим прямо сейчас, по ролям

Исследователям и разработчикам в РФ: если вы работаете с графами знаний на российских данных (графы товаров, социальных сетей, корпоративных баз знаний), DSA снимает порог входа. Pytorch оптимизация памяти через mmap позволяет запускать эксперименты на бюджетной видеокарте без аренды облачных GPU.

Авторам Дзена и контент-специалистам: тема «как обучать большие модели на слабом железе» хорошо работает в технических каналах. Конкретный расчёт (15 ГБ превращаются в пару мегабайт) даёт материал для наглядного поста.

Предпринимателям и стартапам: если ваш продукт использует рекомендательные системы или графы знаний, DSA сокращает затраты на инфраструктуру. Вместо аренды сервера с 80 ГБ видеопамяти можно начать с одной потребительской карты.

Для тех, кто работает с графовыми нейросетями на ограниченном бюджете, DSA даёт конкретный и воспроизводимый способ обойти ограничение памяти. Попробуйте заменить оптимизатор в своём пайплайне (конвейере обработки данных) и замерьте разницу: если расчёт из статьи верен для ваших данных, экономия на железе окупит час потраченного времени.