PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2104.05158v7cs.DCcs.AIcs.LGcs.PF

TL;DR

Large-scale DLRM training must handle trillion-parameter embedding models while meeting production throughput demands, but existing systems face scaling, memory, and communication constraints. The paper presents Neo, a software-hardware co-designed system combining 4D parallelism, optimized embedding execution, and ZionEX hardware. Neo trains 12-trillion-parameter models at up to 1.7 MQPS and achieves up to 40× faster production training than existing solutions.

  • Problem

    Existing DLRM training systems struggle to scale to trillion-parameter models while maintaining throughput and avoiding stale asynchronous updates.

  • Method

    Neo combines cost-aware 4D parallelism, optimized embedding operators, synchronous execution, pipelining, and the co-designed ZionEX platform.

  • Results

    Neo processes up to 1.7 MQPS when training 12-trillion-parameter DLRMs and achieves up to 40× speedup over existing production solutions.

  • Takeaways & Limitations

    The co-designed software and hardware solution supports production training of DLRMs at trillion-parameter scale and can scale toward data-center-scale clusters.

  • Takeaways & Limitations

    Training model-F requires distributing massive embedding tables across nodes and relies on UVM with HBM as a cache; the prototype cluster has 4TB HBM and 24TB DRAM.

Abstract

from arXiv · show

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.

1 INTRODUCTION

DLRMs are infrastructure-intensive because they combine compute-heavy components with thousands of data-intensive embedding operators, yet existing DNN-oriented systems scale poorly for trillion-parameter models. Neo addresses these limitations through 4D parallelism, optimized embedding computation, and the ZionEX hardware platform, achieving up to 40× faster production training.

  • DLRMs are the single largest AI application in data-center infrastructure demand at Meta.
  • Existing DNN-oriented solutions provide suboptimal performance and limited scalability for DLRMs.
  • Data parallelism cannot support DLRMs with up to trillions of parameters because every device must replicate the entire model.
  • DLRM embedding operators create major runtime and memory overheads by requiring thousands of CUDA kernels and up to terabytes of aggregated GPU memory.
  • Neo combines table-wise, row-wise, column-wise, and data parallelism with optimized embedding computation and ZionEX hardware.
  • 40× speedup was achieved for production DLRM training with models containing 12 trillion parameters.

2 BACKGROUND

DLRM offline training must process millions of queries per second and combines dense computation with massive sparse embedding tables. Existing disaggregated parameter-server systems scale these components separately but become impractical for trillion-parameter models, motivating synchronous training with hybrid parallelism.

  • Offline DLRM training is throughput-limited, processing millions of queries per second and tens of petabytes of training data.
  • An embedding operator retrieves rows indexed by nonzero elements of a multi-hot input and combines them through element-wise pooling.
  • Disaggregated parameter-server systems replicate dense parameters while partitioning embedding tables across parameter servers.
  • Parameter-server systems offer separately scalable components and fungible resources for varied DLRM configurations.
  • Trillion-parameter models make parameter-server scaling intractable and can degrade accuracy through stale asynchronous updates.
  • Neo uses synchronous training with 4D parallelism for embeddings, data parallelism for dense operators, pipelining, and AlltoAll communication.

3 OVERVIEW

Neo separates dense and embedding computation strategies, optimizes embedding execution and memory use, and deploys these techniques on the co-designed ZionEX platform. Its end-to-end design also treats data movement as a potential throughput bottleneck.

  • Neo applies data parallelism to compute-intensive DNN layers and 4D parallelism to memory-intensive embedding operators.
  • Neo’s embedding implementation uses hybrid kernel fusion, software-managed caching, and quality-preserving compression.
  • ZionEX is co-designed with Neo’s 4D parallelism to optimize inter-node communication for DLRM training.
  • Data I/O must support fast non-blocking host-to-device transfers and overlap movement with training through buffering or pipelining.

4 4D PARALLELISM

Neo’s 4D parallelism selects among table-, row-, column-, and data-wise schemes to balance memory, communication, and workload across hardware hierarchies. Cost-based placement heuristics and communication-computation pipelining further improve execution efficiency.

  • 4D parallelism: 4D parallelism combines table-wise, row-wise, column-wise, and data parallelism to balance GPU workloads and communication costs.
  • Sharding schemes: Table-wise parallelism minimizes communication but cannot handle tables exceeding one GPU’s memory and may suffer load imbalance.
  • Sharding schemes: Row-wise parallelism handles large tables and improves load balance using bucketization and ReduceScatter.
  • Sharding schemes: Column-wise parallelism enables finer-grained partitioning for large embedding dimensions but replicates input indices and increases their payload.
  • Sharding schemes: Data parallelism is preferable for small embedding tables because it avoids forward-pass communication.
  • Hierarchical parallelism: Hierarchical partitioning improves hardware locality and reduces inter-node communication by exploiting faster intra-node GPU interconnects.
  • Parallelization algorithms: Placement costs combine communication overhead and load imbalance, while polynomial-time greedy and largest-differencing heuristics seek balanced assignments.
  • Pipelining: Neo pipelines inter-batch data movement and overlaps communication with computation to improve GPU utilization.

5 EMBEDDING OPTIMIZATIONS

Neo addresses embedding operators’ kernel-launch and memory-scaling challenges through fusion, multidimensional parallelism, and hierarchical memory management. These techniques support embedding tables with up to trillions of parameters while improving operator performance.

  • Embedding operators launch thousands of GPU kernels per iteration and may contain billions of parameters that exceed one GPU’s memory.
  • Neo combines table-wise, row-wise, column-wise, and data parallelism, recursively applying these dimensions across hardware hierarchy levels.This strategy balances embedding workloads while reducing communication costs.
  • Neo fuses multiple embedding lookups and combines backward propagation with sparse optimizer updates in CUDA kernels.Gradients for the same embedding rows are sorted and aggregated within CUDA thread blocks to preserve correctness.
  • 7×: hybrid fusion improves embedding-computation performance compared with a native implementation.The technique also avoids embedding-gradient allocation and reduces device-memory accesses through shared memory.
  • Neo uses HBM, DRAM, and SSDs as a software-managed memory hierarchy for embedding tables that cannot fit on one GPU.Faster memory levels serve as caches for subsequent layers, while compression techniques further reduce memory requirements.

6 ZIONEX: HARDWARE PLATFORM DESIGN

ZionEX redesigns the earlier Zion platform for scalable distributed DLRM training by adding dedicated GPU-connected networking and retaining modular deployment features. The design addresses Zion’s CPU-mediated communication and scale-out limitations.

  • 6.1 Previous Platform: Zion: Zion provided a heterogeneous single-node design with GPUs for compute-heavy layers and CPUs using DRAM for large embedding operators.Its node included eight CPU sockets, 1.5 TB memory, eight GPUs, and eight NICs.
  • 6.1 Previous Platform: Zion: Zion’s CPU-GPU workload balance required fine-grained partitioning and elaborate pipelining, while heterogeneous execution added transfers and inter-socket communication overheads.
  • 6.1 Previous Platform: Zion: Zion’s CPU-attached NICs forced CPU intervention and extra GPU-CPU transfers, while shared data-center networking introduced congestion and protocol overheads.
  • 6.1 Previous Platform: Zion: Zion was difficult to scale across nodes and could not fully utilize its hardware as DLRM model sizes increased.
  • 6.2 ZionEX: ZionEX reports 1.2 MQPS for a 638 MFLOPS/sample model, equivalent to 766 TFLOPS/s effective performance.The cited comparison reports Zion at approximately 0.25 MQPS and 63 TFLOPS/s effective performance.
  • 6.2 ZionEX: ZionEX uses dedicated RoCE NICs for GPUs and an isolated backend network supporting RDMA and GPUDirect communication.Its extensible design targets interconnection of many thousands of nodes for data-center-scale training.
  • 6.2 ZionEX: ZionEX retains the OAM form factor, modular design, and flexible intra-node accelerator fabric while improving network capabilities.The platform also targets compatibility with existing data-center infrastructure through Open Rack specifications and open standards.

7 IMPLEMENTATION

The implementation combines a PyTorch-based training stack with optimized data ingestion and collective communication for end-to-end DLRM throughput. It restructures embedding inputs and provides sharding-aware GPU processing.

  • 7 IMPLEMENTATION: The training stack uses PyTorch, ATen CUDA operators, and DistributedDataParallel for parameter replication, gradient synchronization, and overlapped backpropagation.
  • 7.1 Data ingestion: Unoptimized data ingestion can create significant latency and disrupt pipelining because DLRMs process orders of magnitude more data than typical DNNs.
  • 7.1 Data ingestion: The previous ingestion format produced roughly a thousand input tensors per iteration for DLRMs with hundreds of embedding tables.This created substantial CPU-GPU transfer overhead and was a key bottleneck for Zion.
  • 7.1 Data ingestion: A combined format uses lengths and concatenated table inputs to consolidate CPU-GPU transfers and feed embedding kernels without layout transformations.
  • 7.1 Data ingestion: Sharding-aware distribution uses AlltoAll communication and custom GPU kernels for permutation, bucketization, and replication across table-wise, row-wise, and column-wise schemes.
  • 7 IMPLEMENTATION: NCCL provides the primary collective communication backend, with PyTorch’s NCCL process group extended for Alltoall and Alltoallv.

8 EVALUATION

The evaluation measures Neo on production DLRMs across training quality, throughput scaling, optimization effects, and model capacity. Neo reaches million-query-per-second throughput on models up to 12 trillion parameters, while exposing remaining communication and scaling limits.

  • Evaluation setup: Neo is evaluated on three production DLRMs covering click-through-rate prediction, ranking, and engagement, using a 16-node ZionEX cluster.The experimental platform uses 8 NVIDIA A100 GPUs per node; the supplied configuration passage identifies the cluster as 16 nodes.
  • End-to-end training: 1.2 MQPS on 128 GPUs delivers a 40× speedup over the previous distributed CPU asynchronous training platform.The comparison uses 45 parameter servers and 15 trainers; the prior platform could not scale further without hurting training quality.
  • Scaling performance: Scaling efficiency reaches around 50% for model-A and around 75% for model-I at larger node counts.Model-A has greater exposed AlltoAll latency because more tables increase payload and mixed dimensions complicate load balancing.
  • Scaling performance: Pipelining hides CPU-to-GPU transfer completely and reduces exposed communication latency below serialized AlltoAll plus AllReduce latency.This demonstrates overlap between communication and computation in model-A.
  • Scaling performance: AlltoAll and AllReduce latencies increase with node count, with AlltoAll affecting exposed latency directly and AllReduce potentially becoming a bottleneck at higher scale.AllReduce remains mostly hidden through 16 nodes but can consume backward-pass slack as scale increases or computation accelerates.
  • Training throughput optimizations: 87% higher training throughput results from combining load balancing, reduced-precision collectives, and increasing global batch size from 64K to 256K.FP16 forward and BF16 backward AlltoAll provide almost 30% speedup without training-quality loss; further large-batch experimentation is warranted.
  • Model capacity: 1.7 MQPS is achieved for model-F with 12 trillion parameters after reducing memory footprint from 96 TB to 24 TB and using hierarchical memory.Row-wise sparse AdaGrad, FP16 embeddings, row-wise sharding, bucketized AlltoAll, ReduceScatter, and UVM-backed HBM caching enable the result.

9 RELATED WORK

Related systems address large-model memory, sparse-data training, communication scheduling, locality, and quantization. Neo’s context is a broader co-designed approach to DLRM training rather than a single communication or memory optimization.

  • Large-model systems: DeepSpeed shards parameters, gradients, and optimizer states across nodes, reconstructing required states through checkpoint partitioning and rematerialization.These techniques target drastic reductions in memory usage for extremely large models.
  • Sparse and online training: XDL targets industry-scale sparse-data training using hierarchical sample compression, workflow pipelining, zero copy, and CPU binding.Kraken instead focuses on online training with decoupled key-value fetching and embedding plus domain-informed cache eviction.
  • Communication systems: BytePS and ByteScheduler improve parameter exchange by using idle CPU and network resources with communication scheduling.Other work exploits datacenter locality and forms optimized aggregation routes.
  • Communication systems: Quantization-based methods reduce communication overhead by reducing communication volume.These approaches are identified as an alternative line of communication optimization.

10 CONCLUSION

DLRMs are widely deployed and impose atypical, rapidly growing demands on training infrastructure. The paper concludes that co-designing Neo with ZionEX supports trillion-parameter production training while pointing toward model and algorithm changes for improved scaling.

  • Conclusion: DLRMs are widely used across internet companies and can be the largest AI application in data-center infrastructure demand.Their requirements differ from other deep learning models while their scale continues to grow.
  • Conclusion: Neo and ZionEX enable trillion-parameter models while attaining 40× faster total training time for production recommendation models.Neo contributes 4D parallelism, high-performance embedding kernels, hybrid kernel fusion, and hierarchical memory management; ZionEX supports scaling toward thousands of nodes.
  • Conclusion: Future work explores model architectures and algorithms that reduce global AlltoAll communication to improve scaling efficiency.The solution has been deployed in production, and these directions aim to extend large-scale training capability.

Compute Benchmarks

The paper uses operator-level benchmarks and workload-aware communication tests to characterize compute, memory, and collective performance on candidate hardware. These measurements support roofline-based performance bounds and expose differences between low-level and application-representative benchmarking.

  • Benchmark methodology: PARAM bench evaluates representative compute, memory, and communication problem sizes and shapes on candidate hardware platforms.The benchmark suite was developed and open sourced to measure throughput and latency across these resource domains.
  • Compute benchmarks: GEMM benchmarks compare V100 FP16 with A100 FP16 and BF16, while MLP benchmarks compare V100 FP32 with A100 FP32 and TF32.The GEMM benchmark uses cuBLAS GemmEx with configurable problem sizes and precision choices.
  • Compute benchmarks: MLP benchmarking covers configurable batch sizes, layer dimensions, 20-layer networks, forward and backward passes, SGD updates, and FP16, BF16, TF32, and FP32 precision.The implementation uses cuBLAS, cuDNN, and custom CUDA kernels for the specified components.
  • Communication benchmarks: Standard NCCL and OSU MPI tests miss exact workload message sizes and operation sequences and are limited to one communication library.PARAM comms addresses these gaps through common abstractions and PyTorch Process Group interfaces across libraries.
  • Embedding benchmarks: Embedding benchmarks measure forward and backward-plus-optimizer bandwidth for FP32 and FP16 kernels on V100 and A100 GPUs.The setup uses randomized 40 MB data to flush the A100 L2 cache and evaluates one million rows across 64 tables.
  • Communication benchmarks: PARAM comms supports bench mode for fixed message sizes and replay mode for traces that reproduce exact collective-call sizes.This provides both low-level hardware testing and workload-mimicking communication evaluation.
  • Roofline analysis: The roofline model decomposes DLRM training into MLP, embedding, AlltoAll, interaction and Top MLP, and AllReduce components.Operator benchmarks estimate each component’s achievable performance, including up to 850 GB/s HBM bandwidth on V100 and 1300 GB/s on A100.
  • Communication benchmarks: On 128 GPUs, AlltoAll reaches 7 GB/s for 256 MB while AllReduce reaches 60 GB/s for 256 MB.AllReduce achieves higher effective bandwidth by using both scale-out and NVLink bandwidths.
Loading 2104.05158v7…