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

TL;DR

The need for a shift to an algorithm-hardware codesign to enable decentralized parallel SGD algorithms in real-world PIM systems, which significantly reduces the communication cost and improves scalability is highlighted.

Abstract

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.

Figures (13)

• Figure 1: High-level system organization of an UPMEM PIM-enabled system and the hardware architecture of an UPMEM PIM chip.
• Figure 2: Per global epoch comparison of measured throughput (a) and total data movement (b) for all studied algorithms (MA-SGD, GA-SGD, and ADMM) on the UPMEM PIM system for the Criteo dataset.
• Figure 3: High-level workflow for distributed optimization algorithms on the UPMEM PIM system.
• Figure 4: Per global epoch training time breakdown into Comm./Sync. Para. Server, PIM Comp., and PIM data movement time for LR (a) and SVM (b).
• Figure 5: Comparison of various models (LR (a) and SVM (b)), algorithms (MA-SGD, GA-SGD, ADMM, and mini-batch SGD), and architectures (PIM, CPU, and GPU). We study the test accuracy (at the last global epoch) and total training time (10 global epochs).