[Project Showcase] Disk Sparse Adam (DSA) — A disk-backed SparseAdam for large embeddings and GNNs

Memory Format
1

Hi PyTorch Community!

I wanted to share an open-source library I’ve developed to solve a common memory bottleneck when training large-scale sparse models: Disk Sparse Adam (DSA).

:red_circle: The Problem: Memory Overheads in Large Sparse Embeddings

When training massive embedding tables (e.g., Knowledge Graph Embeddings or Graph Neural Networks with millions/billions of nodes), standard torch.optim.SparseAdam requires storing first and second momentum states ($m$ and $v$) for every single parameter.

For multi-million-entity tables, these optimizer states quickly overwhelm GPU VRAM or system RAM, leading to CUDA out of memory (OOM) or system crashes—especially on consumer-grade hardware or single-GPU setups.

:green_circle: The Solution: Out-of-Core Sparse Adam (DSA)

Disk Sparse Adam (DSA) is an out-of-core sparse optimizer that offloads SparseAdam momentum states to disk using efficient memory mapping (mmap).

Instead of keeping gigabytes of momentum states in active RAM/VRAM, DSA streams and updates only the active sparse slices required for the current mini-batch, maintaining asynchronous sparse updates with near-zero VRAM overhead for optimizer states.

:sparkles: Key Features

  • Near-Zero VRAM/RAM Footprint for Optimizer States: Offloads $m$ and $v$ tensors to disk via memory mapping.
  • Consumer GPU Friendly: Enables training massive embeddings (10M+ entities) on single consumer GPUs (e.g., RTX 3090/4090 or Google Colab).
  • Drop-in PyTorch Integration: Designed to easily replace standard sparse optimization workflows in PyTorch.
  • Target Workloads: Knowledge Graph Embeddings (KGE), Large-Scale GNN Node Embeddings, Recommendation Systems, and Large Sparse Lookup Tables.

:laptop: Quick Usage Example

import torch
from dsa import DiskSparseAdam

# Massive sparse embedding table (e.g., 10 million entities)
num_embeddings = 10_000_000
embedding_dim = 128
embedding = torch.nn.EmbeddingBag(num_embeddings, embedding_dim, sparse=True)

# Initialize DSA optimizer offloading states to disk
optimizer = DiskSparseAdam(
    embedding.parameters(), 
    lr=0.001, 
    state_dir="./opt_state_cache"
)

# Standard training loop
for batch_indices in dataloader:
    optimizer.zero_grad()
    output = embedding(batch_indices)
    loss = compute_loss(output)
    loss.backward()
    optimizer.step()