Software-hardware co-design for fast and scalable training of deep learning recommendation models

By Dheevatsa Mudigere, Yuchen Hao, Jianyu Huang, Andrew Tulloch, Srinivas Sridharan, Xing Liu, Mustafa Ozdal, Jade Nie, Jongsoo Park, Liangchen Luo and 42 more. 2021

209 citations recorded by the source corpus.

Deep learning recommendation models (DLRMs) have been used across many business-critical services at Meta and are the single largest AI application in terms of infrastructure demand in its data-centers. In this paper, we present Neo, a software-hardware co-designed system for high-performance distributed training of large-scale DLRMs. Neo employs a novel 4D parallelism strategy that combines table-wise, row-wise, column-wise, and data parallelism for training massive embedding operators in DLRMs. In addition, Neo enables extremely high-performance and memory-efficient embedding computations using a variety of critical systems optimizations, including hybrid kernel fusion, software-managed caching, and quality-preserving compression. Finally, Neo is paired with ZionEX, a new hardware platform co-designed with Neo's 4D parallelism for optimizing communications for large-scale DLRM training. Our evaluation on 128 GPUs using 16 ZionEX nodes shows that Neo outperforms existing systems by up to 40× for training 12-trillion-parameter DLRM models deployed in production.

Subfields: Artificial Intelligence, Computer Networks and Communications, Information Systems, Evaluation & Benchmarks, Recommender Systems

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Networks and Communications

In Information Systems

In Evaluation & Benchmarks

In Recommender Systems

Citations in the corpus

Sustainable AI: Environmental Implications, Challenges and Opportunities (2021, background): “The strong performance scaling demand for ML motivates a variety of scale-out solutions [11], [18] by leveraging parallelism at scale with a massive collection of training accelerators.”
TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings (2023, methodology): “Neo from Facebook (Meta) [37] trains embedding tables with up to 12T parameters using a 128-GPU system.”
DHEN: A Deep and Hierarchical Ensemble Network for Large-Scale Click-Through Rate Prediction (2022, methodology): “To allow efficient training of DHEN, we take advantage of the ZionEX fully synchronous training system detailed in [27].”