提出新方法降低推荐系统嵌入表通信开销,提升训练速度。
Stochastic Communication Avoidance for Recommendation Systems
- 设计理论框架分析分布式系统中查表通信成本
- 在Criteo数据集上实现6倍于基线的训练吞吐量提升
- 适用于从移动端到数据中心的各种推荐场景
基于神经网络的推荐系统部署的一大瓶颈是嵌入表带来的巨大内存开销。尽管硬件加速器具备更快的片上内存访问和更强算力,但模型中的大尺寸嵌入表常无法容纳在加速器受限的内存中。现有内存优化与并行方法未能有效解决加速器上大嵌入表的内存与通信开销问题,导致多数模型仍在CPU上训练,而加速器部署受制于设备间通信瓶颈与主存查找延迟。本文提出一个理论框架,用于分析任意使用查表的分布式系统的通信成本,并据此设计算法,在内存、计算与通信约束下最大化吞吐量。我们证明该方法在不同数据分布与内存条件下均具强理论性能,适用于从移动端联邦学习到超大规模数据中心计算等多种场景。我们在PyTorch中实现该框架,在Criteo Terabytes数据集上相较基线实现最高6倍的训练吞吐量提升。
原文摘要 · Abstract (English)
One of the major bottlenecks for efficient deployment of neural network based recommendation systems is the memory footprint of their embedding tables. Although many neural network based recommendation systems could benefit from the faster on-chip memory access and increased computational power of hardware accelerators, the large embedding tables in these models often cannot fit on the constrained memory of accelerators. Despite the pervasiveness of these models, prior methods in memory optimization and parallelism fail to address the memory and communication costs of large embedding tables on accelerators. As a result, the majority of models are trained on CPUs, while current implementations of accelerators are hindered by issues such as bottlenecks in inter-device communication and main memory lookups. In this paper, we propose a theoretical framework that analyses the communication costs of arbitrary distributed systems that use lookup tables. We use this framework to propose algorithms that maximize throughput subject to memory, computation, and communication constraints. Furthermore, we demonstrate that our method achieves strong theoretical performance across dataset distributions and memory constraints, applicable to a wide range of use cases from mobile federated learning to warehouse-scale computation. We implement our framework and algorithms in PyTorch and achieve up to 6x increases in training throughput on GPU systems over baselines, on the Criteo Terabytes dataset.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。