Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models
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.
Cite
@article{arxiv.2104.05158,
title = {Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models},
author = {Dheevatsa Mudigere and Yuchen Hao and Jianyu Huang and Zhihao Jia and Andrew Tulloch and Srinivas Sridharan and Xing Liu and Mustafa Ozdal and Jade Nie and Jongsoo Park and Liang Luo and Jie Amy Yang and Leon Gao and Dmytro Ivchenko and Aarti Basant and Yuxi Hu and Jiyan Yang and Ehsan K. Ardestani and Xiaodong Wang and Rakesh Komuravelli and Ching-Hsiang Chu and Serhat Yilmaz and Huayu Li and Jiyuan Qian and Zhuobo Feng and Yinbin Ma and Junjie Yang and Ellie Wen and Hong Li and Lin Yang and Chonglin Sun and Whitney Zhao and Dimitry Melts and Krishna Dhulipala and KR Kishore and Tyler Graf and Assaf Eisenman and Kiran Kumar Matam and Adi Gangidi and Guoqiang Jerry Chen and Manoj Krishnan and Avinash Nayak and Krishnakumar Nair and Bharath Muthiah and Mahmoud khorashadi and Pallab Bhattacharya and Petr Lapukhov and Maxim Naumov and Ajit Mathews and Lin Qiao and Mikhail Smelyanskiy and Bill Jia and Vijay Rao},
journal= {arXiv preprint arXiv:2104.05158},
year = {2023}
}