PaperlayerView exact revision on arXiv

Source-linked AI summary

Persia: An Open, Hybrid System Scaling Deep Learning-based Recommenders up to 100 Trillion Parameters

Xiangru Lian, Binhang Yuan, Xuefeng Zhu, Yulong Wang, Yongjun He, Honghuan Wu, Lei Sun, Haodong Lyu, Chengjun Liu, Xing Dong, Yiqiao Liao, Mingnan Luo, Congfei Zhang, Jingru Xie, Haonan Li, Lei Chen, Renjie Huang, Jianying Lin, Chengchun Shu, Xuezhong Qiu, Zhishan Liu, Dongying Kong, Lei Yuan, Hai Yu, Sen Yang, Ce Zhang, Ji Liu

arXiv:2111.05897v3cs.LGcs.DC

TL;DR

Training recommender models at 100-trillion-parameter scale is challenging because embedding layers are memory-intensive while neural networks are computation-intensive. Persia combines asynchronous embedding updates with synchronous dense-network training in a co-designed distributed system, reporting up to 7.12× speedup over alternative state-of-the-art approaches. The system is open source and targets training at this scale.

  • Problem

    At 100-trillion-parameter scale, recommender training must handle memory-intensive embeddings and computation-intensive neural networks, while synchronous and asynchronous approaches face efficiency and accuracy trade-offs.

    Source: §4.2.4 · p. 8, §7.1 · p. 12

  • Method

    Persia co-designs a distributed system with hybrid training: asynchronous updates for embeddings and synchronous training for the dense neural network.

    Source: §1 · p. 2

  • Results

    Up to 7.12× speedup over alternative state-of-the-art approaches is reported, alongside evaluation on public benchmarks and Kwai production tasks.

    Source: §1 · p. 2, §8 · p. 12

  • Takeaways & Limitations

    Persia is an open-source system designed to support recommender training at the scale of 100 trillion parameters.

    Source: §1 · p. 2, §8 · p. 12

    Embedding workers have no fault-recovery scheme, and a failure abandons their local ID-feature hash-map buffer.

    Source: §4.2.4 · p. 8

Abstract

from arXiv · show

Deep learning based models have dominated the current landscape of production recommender systems. Furthermore, recent years have witnessed an exponential growth of the model scale--from Google's 2016 model with 1 billion parameters to the latest Facebook's model with 12 trillion parameters. Significant quality boost has come with each jump of the model capacity, which makes us believe the era of 100 trillion parameters is around the corner. However, the training of such models is challenging even within industrial scale data centers. This difficulty is inherited from the staggering heterogeneity of the training computation--the model's embedding layer could include more than 99.99% of the total model size, which is extremely memory-intensive; while the rest neural network is increasingly computation-intensive. To support the training of such huge models, an efficient distributed training system is in urgent need. In this paper, we resolve this challenge by careful co-design of both the optimization algorithm and the distributed system architecture. Specifically, in order to ensure both the training efficiency and the training accuracy, we design a novel hybrid training algorithm, where the embedding layer and the dense neural network are handled by different synchronization mechanisms; then we build a system called Persia (short for parallel recommendation training system with hybrid acceleration) to support this hybrid training algorithm. Both theoretical demonstration and empirical study up to 100 trillion parameters have conducted to justified the system design and implementation of Persia. We make Persia publicly available (at https://github.com/PersiaML/Persia) so that anyone would be able to easily train a recommender model at the scale of 100 trillion parameters.

1 INTRODUCTION

Recommender models are growing toward 100 trillion parameters, but their memory-heavy embedding layers and computation-heavy neural networks make training at that scale challenging. Persia co-designs a hybrid training algorithm and distributed system, reporting up to 7.12× speedup over alternative state-of-the-art approaches.

1 INTRODUCTION

  • The path toward 100 trillion parameters is difficult because embedding layers dominate memory while the remaining neural network is increasingly computation-intensive.At this scale, storing the model requires at least 200TB in fp16, and communication among hundreds of machines can become a bottleneck.Source: §1 · pp. 1–2
  • Existing systems use either synchronous updates, whose communication overhead limits efficiency, or asynchronous updates, which can substantially reduce accuracy at scale.The paper notes that even a 0.1% accuracy drop can noticeably affect revenue in production recommender systems.Source: §1 · p. 2
  • Persia pairs asynchronous embedding-layer training to improve sample throughput with synchronous neural-network training to preserve statistical efficiency.The authors also analyze convergence theoretically and relate recommender-model characteristics to convergence.Source: §1 · p. 2
  • Persia co-designs its distributed system to manage hybrid CPU/GPU computation, optimize memory and communication, and tolerate training failures.The system is open source, with configuration guidance for reproducing the 100-trillion-parameter test on cloud infrastructure.Source: §1 · p. 2
  • Up to 7.12× speedup over alternative state-of-the-art approaches is reported for Persia, evaluated on public benchmarks and Kwai production tasks.The paper also reports scalability experiments up to 100 trillion parameters on public clouds.Source: §1 · pp. 2–3

2 PRELIMINARIES

The preliminaries frame recommender training as a combination of sparse ID-feature embeddings and a dense neural network. This imbalance motivates distributed architectures that separate storage and computation, while Persia's motivation is to combine the benefits of synchronous and asynchronous updates.

2 PRELIMINARIES

  • The preliminaries introduce a typical deep learning recommender and outline the distributed training problem and existing architectures.The example is a DLRM-style recommender.Source: §2 · p. 3

2.1 Problem Formalization

  • The formalization separates ID-type features and their embedding parameters from non-ID features and the dense neural-network parameters.Embedding parameters can be 10^7 times larger in size, while neural-network computation can be 10^7 times greater; this imbalance motivates specialized training systems.Source: §2.1 · p. 3

2.2 Anatomy of Existing Architectures

  • Existing architectures distribute embedding storage across CPU parameter-server nodes and assign intensive computation to GPUs, but systems generally use either synchronous or asynchronous updates.Persia is motivated by combining the hardware efficiency of asynchronous updates with the accuracy advantages of synchronous updates for large models.Source: §2.2 · pp. 3–4

3 HYBRID TRAINING ALGORITHM

Persia pairs asynchronous embedding updates with synchronous dense-network updates to balance hardware and statistical efficiency. Its algorithm separates embedding lookups and gradient updates from mini-batch neural-network training.

3.1 Why Hybrid Algorithm?

  • The algorithm addresses the trade-off between synchronous communication overhead and asynchronous staleness that can impair scalability, convergence, and generalization.The section motivates separate update mechanisms for embeddings and the dense network.Source: §3.1 · p. 4
  • The hybrid algorithm updates embeddings asynchronously while keeping dense-network training synchronous, aiming for asynchronous-like hardware efficiency without sacrificing statistical efficiency.System optimizations can hide dense-parameter synchronization within backward propagation.Source: §3.1 · p. 4

3.2 Asynchronous updating for wemb

  • Embedding forward tasks retrieve vectors and send them onward, while backward tasks send their gradients to parameter storage without synchronization locks.The embedding update algorithm runs these forward and backward tasks independently.Source: §3.2 · p. 4

3.3 Synchronous updating for wnn

  • Dense-network training uses mini-batch SGD with synchronization locks, and concatenates embedding activations into the neural network’s input.The algorithm otherwise resembles standard distributed neural-network training.Source: §3.3 · p. 5

4 SYSTEM DESIGN AND IMPLEMENTATION

Persia separates embedding and neural-network training across CPU and GPU resources, using different communication paradigms for each. It adds buffering, memory-management, communication, and fault-tolerance mechanisms to support the hybrid algorithm.

4.1 System Design

  • Persia assigns embedding storage and updates to a CPU-based parameter server and embedding workers, while GPU-based neural-network workers synchronize gradients with AllReduce.The architecture also includes a data loader, and its workflow dispatches samples across these modules.Source: §4.1 · p. 6
  • Persia’s system goals include bridging asynchronous and synchronous computation, using heterogeneous resources efficiently, and tolerating failures.The hybrid workflow requires coordination of embeddings and gradients across a large-scale cluster.Source: §4.1 · p. 6

4.2.1 Fill the Async/Sync Gap.

  • Local sample-ID-keyed buffers let embedding and neural-network workers coordinate embedding vectors, non-ID features, labels, and gradients across asynchronous and synchronous work.The NN worker matches arriving embedding vectors with stored sample inputs; the embedding worker maps sample IDs back to ID features.Source: §4.2.1 · pp. 6–7

4.2.2 Persia Memory Management.

  • The embedding parameter server uses an LRU cache implemented with a hash map and array-based linked list, avoiding fragmented allocations and simplifying checkpoint serialization.Array indices replace memory-address pointers; multiple threads manage subsets of the cache.Source: §4.2.2 · p. 7

4.2.3 Communication Optimization.

  • For NN-worker synchronization, Persia uses Bagua’s centralized synchronous full-precision communication by default; its RPC uses zero-copy tensor serialization.Embedding parameters are shuffled within feature groups and distributed evenly across server nodes to reduce access congestion and balance workload.Source: §4.2.3 · pp. 7–8
  • Embedding-to-NN communication compresses indices losslessly and embedding values with a discreet fp32-to-fp16 mapping.The index representation uses uint16 for batch sizes up to 65535; the value mapping scales each vector block by κ times its L∞ norm before conversion.Source: §4.2.3 · p. 8

4.2.4 Fault Tolerance.

  • Fault handling reflects different convergence sensitivities: embedding-worker failures abandon local buffers, while NN-worker failures trigger reload of the latest synchronized-model checkpoint.Embedding parameter shards are periodically checkpointed; the data loader relies on its underlying storage system’s recovery scheme.Source: §4.2.4 · p. 8

5 THEORETICAL ANALYSIS

Persia’s hybrid algorithm updates embeddings asynchronously and the dense network synchronously. Under the stated assumptions, it benefits from asynchrony while retaining a convergence rate similar to synchronous training in typical recommendation settings.

Theoretical analysis

  • The hybrid update uses an earlier iterate for embeddings but the up-to-date dense-network value, reflecting asynchronous embedding updates and synchronous dense-network updates.The paper characterizes this staleness cost as minor compared with the system-efficiency gain.Source: §5 · p. 9
  • The convergence analysis assumes a global minimum exists, stochastic gradients have bounded variance and Lipschitz gradients, and sparse-update delays are bounded.In Persia, the staleness bound τ is less than 5 in most cases.Source: §5 · p. 9
  • The theorem’s ID-frequency constant α bounds the probability that a sample contains any given ID; α = 1/N_emb when IDs are equally likely across samples.The frequency assumption reflects an intrinsic property of recommendation systems.Source: §5 · p. 9
  • When α ≪ 1, as is usually true in recommendation systems, the staleness term is dominated and the hybrid algorithm has a convergence rate similar to synchronous SGD.At α = 1, its rate matches asynchronous SGD, so the theorem indicates it is no worse than asynchronous SGD in convergence rate.Source: §5 · p. 9

6 EXPERIMENTS

Persia is evaluated on three open-source ad benchmarks and Kwai’s production video-recommendation task. The experiments report faster end-to-end training, hybrid-algorithm convergence close to synchronous training, and scaling to 100 trillion parameters.

6 EXPERIMENTS

  • Persia evaluates Avazu-Ad and Criteo-Ad as open-source advertisement CTR benchmarks and Kwai-Video as a production microvideo recommendation task, using test AUC to evaluate convergence.The open-source tasks use 80/20 train/test splits, while Kwai-Video uses 85/15. Comparisons include XDL and PaddlePaddle baselines.Source: §6 · p. 10
  • Experiments span production and Google Cloud clusters, including a capacity setup with 8 GPU-worker instances, 100 embedding-worker instances, and 30 embedding parameter-server instances.Baseline systems receive the same computation resources as Persia in each individual setting.Source: §6 · p. 10

6.1 End-to-end Evaluation

  • 7.12× faster than synchronous XDL and 8.4× faster than asynchronous XDL on Taobao-Ad; Persia is also 1.71× faster than PaddlePaddle.The authors report similar levels of speedup on Avazu-Ad and Criteo-Ad; XDL did not converge on Kwai-Video within 100 hours, where Persia’s sample throughput was 19.31× XDL’s.Source: §6.1 · p. 10, Figure 6 · p. 10

6.2 Evaluation of Hybrid Algorithm

  • Hybrid training keeps test-AUC gaps below 0.1% on the three open-source benchmarks and below 0.001% on Kwai-Video versus synchronous training.Asynchronous training has a larger 0.5%–1.0% gap, increasing with more aggressive asynchronicity in PaddlePaddle.Source: §6.2 · pp. 10–11, Figure 7 · p. 11, Table 2 · p. 11

6.3 Scalability and Capacity

  • Persia achieves nearly linear speedup with higher throughput than XDL and PaddlePaddle on the three open-source benchmarks.On Kwai-Video, Persia’s throughput is 3.8× that of fully synchronous training; fully asynchronous training is faster but has lower AUC than hybrid training.Source: §6.3 · p. 11
  • At 100 trillion parameters, Persia achieves 2.6× the throughput of fully synchronous training; asynchronous training is 1.2× faster than hybrid training.Throughput remains stable as the Criteo-Syn embedding-table size increases to 100 trillion parameters; the faster asynchronous mode’s convergence efficiency is not assessed by this capacity test.Source: §6.3 · p. 11

7 RELATED WORK

Related work covers recommender systems and distributed deep learning. It describes existing heterogeneous architectures and distributed-learning approaches while identifying a lack of solutions that jointly address recommender-training heterogeneity across algorithms and infrastructure.

7.1 Recommender Systems

  • Deep-learning recommenders map sparse ID features through embeddings before neural components model their interactions, and may incorporate sequential, text, image, audio, video, social, or multi-task features.Source: §7.1 · p. 11
  • Prior recommender systems use heterogeneous resources and mechanisms such as model servers, hierarchical parameter servers, and GPU-side embedding caches.The paper identifies a gap in solutions designed for sparse-versus-dense data and memory-versus-computation constraints over mixed CPU/GPU infrastructure.Source: §7.1 · pp. 11–12

7.2 Distributed Deep Learning

  • Distributed deep-learning work includes parameter-server and AllReduce data parallelism, communication-reduction methods, asynchronous updates, and model-parallel approaches.Some approaches combine strategies, including data and model parallelism.Source: §7.2 · p. 12
  • General-purpose deep-learning systems and their extensions do not account for hybrid-infrastructure heterogeneity, making them difficult to use directly for industrial-scale recommender training on heterogeneous clusters.Source: §7.2 · p. 12

8 CONCLUSION

The paper presents Persia as a co-designed distributed training system and algorithm for recommender models up to 100 trillion parameters. It reports up to 7.12× speedup over state-of-the-art alternatives and releases the system with setup instructions.

8 CONCLUSION

  • Persia co-designs a distributed training algorithm and system to train recommender models at up to 100 trillion parameters.Source: §8 · p. 12
  • The hybrid algorithm uses system relaxations for heterogeneous-cluster utilization while converging as fast as vanilla SGD.Source: §8 · p. 12
  • Persia reports up to 7.12× speedup over alternative state-of-the-art approaches and scales to a 100-trillion-parameter model on Google Cloud Platform.The project is open source and includes setup instructions for Google Cloud Platform.Source: §8 · p. 12

APPENDIX

The appendix derives bounds used to characterize the sparse-weight objective and its convergence. The derivation relies on smoothness and stochastic-gradient assumptions and sufficiently small learning rates.

APPENDIX

  • The proof bounds the sparse-weight Hessian using the lookup structure: unused embedding weights have zero contribution to the sample Hessian.Taking the expectation yields an αL bound and αL-Lipschitz continuity in the sparse weights.Source: APPENDIX · p. 16
  • The convergence derivation uses unbiased stochastic gradients and bounded variance to bound terms in its progress inequality.Source: APPENDIX · pp. 16–17
  • The derivation obtains a final convergence-rate bound by summing (12) from t = 0 to T−1 and choosing sufficiently small γ satisfying 1/4 ≥ γτL_emb; an earlier step requires 1/2 ≥ γτL_emb.Source: APPENDIX · pp. 17–19

Related papers

Loading 2111.05897v3…