# PIM-Opt: Demystifying Distributed Optimization Algorithms on a Real-World Processing-In-Memory System

arXiv: https://arxiv.org/abs/2404.07164 · HTML: https://mlsanity.com/paper/2404.07164/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 68 of its 211 references that are in the map.
- First posted on arXiv 10 Apr 2024.

## Builds on (most cited first, 30 of 68)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Can Decentralized Algorithms Outperform Centralized Algorithms? A Case Study for Decentralized Parallel Stochastic Gradient Descent](https://mlsanity.com/paper/1705.09056/) (2017)
- [Sparsified SGD with Memory](https://mlsanity.com/paper/1809.07599/) (2018)
- [Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning](https://mlsanity.com/paper/1807.06629/) (2018)
- [FINN: A Framework for Fast, Scalable Binarized Neural Network Inference](https://mlsanity.com/paper/1612.07119/) (2016)
- [Towards Theoretically Understanding Why SGD Generalizes Better Than ADAM in Deep Learning](https://mlsanity.com/paper/2010.05627/) (2020)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [On the convergence properties of a K-step averaging stochastic gradient descent algorithm for nonconvex optimization](https://mlsanity.com/paper/1708.01012/) (2017)
- [To Repeat or Not To Repeat: Insights from Scaling LLM under Token-Crisis](https://mlsanity.com/paper/2305.13230/) (2023)
- [Training Transformers with 4-bit Integers](https://mlsanity.com/paper/2306.11987/) (2023)
- [RecNMP: Accelerating Personalized Recommendation with Near-Memory Processing](https://mlsanity.com/paper/1912.12953/) (2019)
- [Communication-Efficient Algorithms for Decentralized and Stochastic Optimization](https://mlsanity.com/paper/1701.03961/) (2017)
- [TensorDIMM: A Practical Near-Memory Processing Architecture for Embeddings and Tensor Operations in Deep Learning](https://mlsanity.com/paper/1908.03072/) (2019)
- [A Survey on Large-scale Machine Learning](https://mlsanity.com/paper/2008.03911/) (2020)
- [Parallel SGD: When does averaging help?](https://mlsanity.com/paper/1606.07365/) (2016)
- [A Modern Primer on Processing in Memory](https://mlsanity.com/paper/2012.03112/) (2020)
- [Google Neural Network Models for Edge Devices: Analyzing and Mitigating Machine Learning Inference Bottlenecks](https://mlsanity.com/paper/2109.14320/) (2021)
- [Processing Data Where It Makes Sense: Enabling In-Memory Computation](https://mlsanity.com/paper/1903.03988/) (2019)
- [lo-fi: distributed fine-tuning without communication](https://mlsanity.com/paper/2210.11948/) (2022)
- [Distributed Learning Systems with First-order Methods](https://mlsanity.com/paper/2104.05245/) (2021)
- [SISA: Set-Centric Instruction Set Architecture for Graph Mining on Processing-in-Memory Systems](https://mlsanity.com/paper/2104.07582/) (2021)
- [Towards Demystifying Serverless Machine Learning Training](https://mlsanity.com/paper/2105.07806/) (2021)
- [Pathfinding Future PIM Architectures by Demystifying a Commercial PIM Technology](https://mlsanity.com/paper/2308.00846/) (2023)
- [Benchmarking Memory-Centric Computing Systems: Analysis of Real Processing-in-Memory Hardware](https://mlsanity.com/paper/2110.01709/) (2021)
- [Memory and Communication Efficient Distributed Stochastic Optimization with Minibatch-Prox](https://mlsanity.com/paper/1702.06269/) (2017)
- [GRIM-Filter: Fast Seed Location Filtering in DNA Read Mapping Using Processing-in-Memory Technologies](https://mlsanity.com/paper/1711.01177/) (2017)

## Built on by (most cited first, 5 of 5)

- [PAPI: Exploiting Dynamic Parallelism in Large Language Model Decoding with a Processing-In-Memory-Enabled Computing System](https://mlsanity.com/paper/2502.15470/) (2025)
- [SwiftRL: Towards Efficient Reinforcement Learning on Real Processing-In-Memory Systems](https://mlsanity.com/paper/2405.03967/) (2024)
- [UpANNS: Enhancing Billion-Scale ANNS Efficiency with Real-World PIM Architecture](https://mlsanity.com/paper/2410.23805/) (2024)
- [Detecting Anomalies in Machine Learning Infrastructure via Hardware Telemetry](https://mlsanity.com/paper/2510.26008/) (2025)
- [DARTH-PUM: A Hybrid Processing-Using-Memory Architecture](https://mlsanity.com/paper/2602.16075/) (2026)

## Abstract (arXiv)

Modern Machine Learning (ML) training on large-scale datasets is a very time-consuming workload. It relies on the optimization algorithm Stochastic Gradient Descent (SGD) due to its effectiveness, simplicity, and generalization performance. Processor-centric architectures (e.g., CPUs, GPUs) commonly used for modern ML training workloads based on SGD are bottlenecked by data movement between the processor and memory units due to the poor data locality in accessing large datasets. As a result, processor-centric architectures suffer from low performance and high energy consumption while executing ML training workloads. Processing-In-Memory (PIM) is a promising solution to alleviate the data movement bottleneck by placing the computation mechanisms inside or near memory. Our goal is to understand the capabilities of popular distributed SGD algorithms on real-world PIM systems to accelerate data-intensive ML training workloads. To this end, we 1) implement several representative centralized parallel SGD algorithms on the real-world UPMEM PIM system, 2) rigorously evaluate these algorithms for ML training on large-scale datasets in terms of performance, accuracy, and scalability, 3) compare to conventional CPU and GPU baselines, and 4) discuss implications for future PIM hardware and highlight the need for a shift to an algorithm-hardware codesign. Our results demonstrate three major findings: 1) The UPMEM PIM system can be a viable alternative to state-of-the-art CPUs and GPUs for many memory-bound ML training workloads, especially when operations and datatypes are natively supported by PIM hardware, 2) it is important to carefully choose the optimization algorithms that best fit PIM, and 3) the UPMEM PIM system does not scale approximately linearly with the number of nodes for many data-intensive ML training workloads. We open source all our code to facilitate future research.