[论文解读] High-performance, Distributed Training of Large-scale Deep Learning Recommendation Models.

Dheevatsa Mudigere, Yuchen Hao+49 more|arXiv (Cornell University)|Apr 12, 2021

Stochastic Gradient Optimization Techniques参考文献 56被引用 18

一句话总结

本论文提出了一套软硬件协同设计的系统，用于大规模深度学习推荐模型（DLRMs）的高性能分布式训练，相比先前系统实现了40倍的加速。该系统结合基于PyTorch的训练堆栈与ZionEX平台，通过优化的数据/模型并行、分层嵌入分片、低精度通信以及高带宽低延迟的网络基础设施，实现了对高达12万亿参数模型的训练。

ABSTRACT

Deep learning recommendation models (DLRMs) are used across many business-critical services at Facebook and are the single largest AI application in terms of infrastructure demand in its data-centers. In this paper we discuss the SW/HW co-designed solution for high-performance distributed training of large-scale DLRMs. We introduce a high-performance scalable software stack based on PyTorch and pair it with the new evolution of Zion platform, namely ZionEX. We demonstrate the capability to train very large DLRMs with up to 12 Trillion parameters and show that we can attain 40X speedup in terms of time to solution over previous systems. We achieve this by (i) designing the ZionEX platform with dedicated scale-out network, provisioned with high bandwidth, optimal topology and efficient transport (ii) implementing an optimized PyTorch-based training stack supporting both model and data parallelism (iii) developing sharding algorithms capable of hierarchical partitioning of the embedding tables along row, column dimensions and load balancing them across multiple workers; (iv) adding high-performance core operators while retaining flexibility to support optimizers with fully deterministic updates (v) leveraging reduced precision communications, multi-level memory hierarchy (HBM+DDR+SSD) and pipelining. Furthermore, we develop and briefly comment on distributed data ingestion and other supporting services that are required for the robust and efficient end-to-end training in production environments.

研究动机与目标

• 应对生产环境中大规模深度学习推荐模型（DLRMs）日益增长的基础设施需求。
• 克服参数量达万亿级别的DLRMs在分布式训练中的性能瓶颈。
• 构建可扩展、高吞吐量的训练系统，在最大化硬件利用率的同时保持确定性优化器更新。
• 通过生产环境中集成的数据摄入与支持服务，实现端到端的训练效率。
• 在不牺牲模型精度或训练稳定性的情况下，显著提升相比先前系统的训练速度。

提出的方法

• 设计并部署ZionEX平台，配备专用横向扩展网络，具备高带宽、最优拓扑结构和高效传输协议。
• 实现基于PyTorch的训练堆栈，原生支持跨分布式工作节点的模型并行与数据并行。
• 开发分层分片算法，沿行与列维度对嵌入表进行分区，并在多个工作节点间实现负载均衡。
• 集成高性能核心算子，全面支持确定性优化器更新，保障训练可复现性。
• 利用低精度通信、多级内存层次（HBM、DDR、SSD）以及通信流水线技术，降低延迟与带宽压力。
• 构建分布式数据摄入流水线与支持服务，确保在生产环境中实现稳健、可扩展且高效的端到端训练。

实验结果

研究问题

• RQ1 如何在生产数据中心环境中实现对参数量高达12万亿的DLRMs分布式训练的可扩展性与高效性？
• RQ2 为实现相比先前训练系统40倍的加速，需要哪些系统级优化？
• RQ3 如何对嵌入表进行分层分区，以在最小化通信开销的同时实现工作节点间的负载均衡？
• RQ4 低精度通信与内存层次优化在多大程度上可提升训练吞吐量，同时不损害训练稳定性？
• RQ5 软硬件协同组件在实现大规模模型的确定性、高性能训练中扮演何种关键作用？

主要发现

• 与先前训练系统相比，该系统在求解时间上实现了40倍的加速。
• 该平台支持训练高达12万亿参数的DLRMs，证明了其可扩展至目前已知最大的推荐模型规模。
• 嵌入表的分层分片有效实现了分布式工作节点间的负载均衡，同时最小化了通信瓶颈。
• 低精度通信与流水线技术显著降低了训练延迟，提升了硬件利用率。
• HBM、DDR与SSD内存层级的集成，实现了对大规模模型状态与梯度的高效处理。
• 系统保持了完全确定性的优化器更新，确保在采用激进优化措施的同时仍具备训练可复现性。