Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Dheevatsa Mudigere, Yuchen Hao, Jianyu Huang, Zhihao Jia, Andrew Tulloch, Srinivas Sridharan, Xing Liu, Mustafa Ozdal, Jade Nie, Jongsoo Park, Liang Luo, Jie Amy Yang, Leon Gao, Dmytro Ivchenko, Aarti Basant, Yuxi Hu, Jiyan Yang, Ehsan K. Ardestani, Xiaodong Wang, Rakesh Komuravelli, Ching-Hsiang Chu, Serhat Yilmaz, Huayu Li, Jiyuan Qian, Zhuobo Feng, Yinbin Ma, Junjie Yang, Ellie Wen, Hong Li, Lin Yang, Chonglin Sun, Whitney Zhao, Dimitry Melts, Krishna Dhulipala, KR Kishore, Tyler Graf, Assaf Eisenman, Kiran Kumar Matam, Adi Gangidi, Guoqiang Jerry Chen, Manoj Krishnan, Avinash Nayak, Krishnakumar Nair, Bharath Muthiah, Mahmoud khorashadi, Pallab Bhattacharya, Petr Lapukhov, Maxim Naumov, Ajit Mathews, Lin Qiao, Mikhail Smelyanskiy, Bill Jia, Vijay Rao

TL;DR

The paper addresses the challenge of training deep learning recommendation models (DLRMs) at trillion-parameter scales with high throughput. It introduces Neo, a software-hardware co-design built on PyTorch and the ZionEX platform, leveraging 4D parallelism, high-performance embedding kernels, and memory-hierarchy-aware optimizations to achieve up to 40x faster training. Key contributions include the 4D embedding partitioning strategy, hybrid kernel fusion for embedding computation, a software cache-based memory hierarchy, and a scalable ZionEX interconnect for distributed training. The results demonstrate end-to-end training of trillion-parameter DLRMs with up to 1.7 MQPS and provide a blueprint for production-scale, data-center-ready AI training for sparse, embedding-heavy workloads.

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.

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

TL;DR

The paper addresses the challenge of training deep learning recommendation models (DLRMs) at trillion-parameter scales with high throughput. It introduces Neo, a software-hardware co-design built on PyTorch and the ZionEX platform, leveraging 4D parallelism, high-performance embedding kernels, and memory-hierarchy-aware optimizations to achieve up to 40x faster training. Key contributions include the 4D embedding partitioning strategy, hybrid kernel fusion for embedding computation, a software cache-based memory hierarchy, and a scalable ZionEX interconnect for distributed training. The results demonstrate end-to-end training of trillion-parameter DLRMs with up to 1.7 MQPS and provide a blueprint for production-scale, data-center-ready AI training for sparse, embedding-heavy workloads.

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.

Paper Structure

This paper contains 37 sections, 2 equations, 18 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Our Approach
  3. 4D parallelism.
  4. High-performance embedding computation.
  5. Hardware platform design.
  6. Results.
  7. Background
  8. Embedding operators.
  9. Parallelization strategies.
  10. Overview
  11. 4D Parallelism
  12. Table-wise parallelism.
  13. Row-wise parallelism.
  14. Column-wise parallelism.
  15. Data parallelism.
  16. ...and 22 more sections

Figures (18)

  • Figure 1: Comparing deep learning models in total amount of compute, in petaflop/s-days (top) ai-and-compute and model capacity (bottom).
  • Figure 2: Workflow of an embedding operator.
  • Figure 3: Disaggregated parameter-server based system
  • Figure 4: Neo overview. Each box in the figure indicates a neural network component, while edges between boxes are tensors shared between different components.
  • Figure 5: Embedding table sharding schemes with different implications on the communication cost, load balancing and memory requirement. Bottom MLP is omitted in this figure for simplicity of illustration.
  • ...and 13 more figures