摘要
现代机器学习 (ML) 在大规模数据集上的训练是一个非常耗时的工作负载。它依赖于优化算法随机梯度下降 (SGD),因为它有效、简单且具有泛化性能。用于基于 SGD 的现代 ML 训练工作负载的以处理器为中心的架构(例如,CPU、GPU)由于访问大型数据集的数据局部性差,而受到处理器和内存单元之间数据移动的瓶颈。因此,以处理器为中心的架构在执行 ML 训练工作负载时会遭受性能低下和能耗高的困扰。内存中处理 (PIM) 是一种很有前景的解决方案,可以通过将计算机制放置在内存内部或附近来缓解数据移动瓶颈。
我们的目标是了解流行的分布式 SGD 算法在真实 PIM 系统上的能力,以加速数据密集型 ML 训练工作负载。为此,我们 1) 在真实 UPMEM PIM 系统上实现了几种具有代表性的集中式并行 SGD 算法,2) 严格评估这些算法在性能、准确性和可扩展性方面对大型数据集的 ML 训练,3) 与传统的 CPU 和 GPU 基线进行比较,以及 4) 讨论对未来 PIM 硬件的影响,并强调需要转向算法-硬件协同设计。
我们的结果表明了三个主要发现:1) UPMEM PIM 系统可以成为许多内存绑定 ML 训练工作负载的最新 CPU 和 GPU 的可行替代方案,尤其是在 PIM 硬件原生支持操作和数据类型时,2) 重要的是要仔细选择最适合 PIM 的优化算法,以及 3) UPMEM PIM 系统在许多数据密集型 ML 训练工作负载中无法随着节点数量的增加而近似线性扩展。我们开源了我们所有的代码,以促进未来的研究。