DSA Saves PyTorch Memory, but Passes the Cost On to the SSD
Sh0ny
9 августа 20262 min read
In short
Disk Sparse Adam allows you to train massive sparse embeddings without consuming gigabytes of RAM and VRAM for the optimizer’s state. However, this is not a speedup: memory is saved through constant data exchange with the disk, and its scope of application is limited to sparse gradients.
With large graph models, the problem is often not the parameters themselves, but the optimizer's memory. For 10 million entities with 128-dimensional vectors, the weight table takes up about 5.12 GB, and standard SparseAdam adds another approximately 10.24 GB for two Adam states.
Disk Sparse Adam (DSA) solves this in a fairly straightforward way: it stores the m and v states not in RAM or VRAM, but in files on disk, mounted via mmap. Since only a small portion of the entities is updated in each sparse batch, the optimizer reads only the necessary indices from disk, updates them, and writes them back.
The practical benefit is clear: a model that previously wouldn’t fit into the memory of a standard workstation now has a chance to run on a single GPU or even in the free Colab environment. This can be useful for graph embeddings, Knowledge Graph Embeddings, and user and product recommendation tables—anywhere where the parameters are massive but only a small portion of them changes at each step.
But the key trade-off here isn’t in the Adam formula, but in where the data is stored. DSA doesn’t make training faster or reduce the total volume of states—it simply swaps expensive memory for disk I/O. The paper presents a synthetic run on 1 million entities with a claimed 0 MB of additional VRAM and 134,212 samples per second, but there is no comparison with standard SparseAdam, so these figures cannot be considered proof of acceleration.
There are significant limitations: a fast NVMe SSD is recommended, and an HDD could become a bottleneck. The optimizer is designed specifically for sparse parameters such as Embedding and EmbeddingBag; for dense convolutional layers or transformers, this approach makes no sense. Furthermore, the usage example requires manually reading weights by index and passing gradients to optimizer.step(), so this is more of a specialized tool for lookup tables than a universal drop-in replacement.
If you’re facing an OOM error due to a massive embedding table, DSA seems like a reasonable way to buy some time and avoid renting a server. But if you have enough memory, you should first measure your I/O speed: are you willing to trade a VRAM shortage for a constant reliance on an SSD?