# Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models

arXiv: https://arxiv.org/abs/2104.05158 · HTML: https://mlsanity.com/paper/2104.05158/

## Key facts

- Cited by 78 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 75 references that are in the map.
- First posted on arXiv 12 Apr 2021.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Xception: Deep Learning with Depthwise Separable Convolutions](https://mlsanity.com/paper/1610.02357/) (2016)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Neural Collaborative Filtering](https://mlsanity.com/paper/1708.05031/) (2017)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training](https://mlsanity.com/paper/1712.01887/) (2017)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [MLPerf Inference Benchmark](https://mlsanity.com/paper/1911.02549/) (2019)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [MLPerf Training Benchmark](https://mlsanity.com/paper/1910.01500/) (2019)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Deep learning with Elastic Averaging SGD](https://mlsanity.com/paper/1412.6651/) (2014)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [The Architectural Implications of Facebook's DNN-based Personalized Recommendation](https://mlsanity.com/paper/1906.03109/) (2019)
- [Scaling Distributed Machine Learning with In-Network Aggregation](https://mlsanity.com/paper/1903.06701/) (2019)
- [Dynamic Tensor Rematerialization](https://mlsanity.com/paper/2006.09616/) (2020)
- [Understanding Training Efficiency of Deep Learning Recommendation Models at Scale](https://mlsanity.com/paper/2011.05497/) (2020)
- [Deep Learning Training in Facebook Data Centers: Design of Scale-up and Scale-out Systems](https://mlsanity.com/paper/2003.09518/) (2020)
- [Deep Learning Inference in Facebook Data Centers: Characterization, Performance Optimizations and Hardware Implications](https://mlsanity.com/paper/1811.09886/) (2018)

## Built on by (most cited first, 30 of 78)

- [Sustainable AI: Environmental Implications, Challenges and Opportunities](https://mlsanity.com/paper/2111.00364/) (2021)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations](https://mlsanity.com/paper/2402.17152/) (2024)
- [Compute Trends Across Three Eras of Machine Learning](https://mlsanity.com/paper/2202.05924/) (2022)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Machine Learning Model Sizes and the Parameter Gap](https://mlsanity.com/paper/2207.02852/) (2022)
- [TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches](https://mlsanity.com/paper/2111.04867/) (2021)
- [Understanding Scaling Laws for Recommendation Models](https://mlsanity.com/paper/2208.08489/) (2022)
- [MoESys: A Distributed and Efficient Mixture-of-Experts Training and Inference System for Internet Services](https://mlsanity.com/paper/2205.10034/) (2022)
- [Understanding Data Storage and Ingestion for Large-Scale Deep Recommendation Model Training](https://mlsanity.com/paper/2108.09373/) (2021)
- [CASSINI: Network-Aware Job Scheduling in Machine Learning Clusters](https://mlsanity.com/paper/2308.00852/) (2023)
- [Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Models](https://mlsanity.com/paper/2110.04478/) (2021)
- [Unified Embedding: Battle-Tested Feature Representations for Web-Scale ML Systems](https://mlsanity.com/paper/2305.12102/) (2023)
- [MAD Max Beyond Single-Node: Enabling Large Machine Learning Model Acceleration on Distributed Systems](https://mlsanity.com/paper/2310.02784/) (2023)
- [Heterogeneous Acceleration Pipeline for Recommendation System Training](https://mlsanity.com/paper/2204.05436/) (2022)
- [Rail-only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameters](https://mlsanity.com/paper/2307.12169/) (2023)
- [DHEN: A Deep and Hierarchical Ensemble Network for Large-Scale Click-Through Rate Prediction](https://mlsanity.com/paper/2203.11014/) (2022)
- [Experimental Analysis of Large-scale Learnable Vector Storage Compression](https://mlsanity.com/paper/2311.15578/) (2023)
- [Training Personalized Recommendation Systems from (GPU) Scratch: Look Forward not Backwards](https://mlsanity.com/paper/2205.04702/) (2022)
- [BagPipe: Accelerating Deep Recommendation Model Training](https://mlsanity.com/paper/2202.12429/) (2022)
- [RecD: Deduplication for End-to-End Deep Learning Recommendation Model Training Infrastructure](https://mlsanity.com/paper/2211.05239/) (2022)
- [MP-Rec: Hardware-Software Co-Design to Enable Multi-Path Recommendation](https://mlsanity.com/paper/2302.10872/) (2023)
- [HammingMesh: A Network Topology for Large-Scale Deep Learning](https://mlsanity.com/paper/2209.01346/) (2022)
- [Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning](https://mlsanity.com/paper/2408.14158/) (2024)
- [DV365: Extremely Long User History Modeling at Instagram](https://mlsanity.com/paper/2506.00450/) (2025)
- [CAFE: Towards Compact, Adaptive, and Fast Embedding for Large-scale Recommendation Models](https://mlsanity.com/paper/2312.03256/) (2023)
- [Impact of RoCE Congestion Control Policies on Distributed Training of DNNs](https://mlsanity.com/paper/2207.10898/) (2022)
- [cedar: Optimized and Unified Machine Learning Input Data Pipelines](https://mlsanity.com/paper/2401.08895/) (2024)
- [Disaggregated Multi-Tower: Topology-aware Modeling Technique for Efficient Large-Scale Recommendation](https://mlsanity.com/paper/2403.00877/) (2024)
- [GPU-based Private Information Retrieval for On-Device Machine Learning Inference](https://mlsanity.com/paper/2301.10904/) (2023)

## Abstract (arXiv)

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.