Two-dimensional Sparse Parallelism for Large Scale Deep Learning Recommendation Model Training

TL;DR

This paper introduces a two-dimensional sparse parallelism method for large-scale deep learning recommendation models, significantly improving training efficiency and scalability on massive GPU clusters while maintaining model accuracy.

Contribution

It proposes a novel 2D sparse parallelism approach combining data and model parallelism, addressing scalability and communication challenges in large DLRMs.

Findings

Achieves nearly linear speedup up to 4,000 GPUs.

Reduces peak memory consumption during training.

Sets new state-of-the-art benchmark for recommendation model training.

Abstract

The increasing complexity of deep learning recommendation models (DLRM) has led to a growing need for large-scale distributed systems that can efficiently train vast amounts of data. In DLRM, the sparse embedding table is a crucial component for managing sparse categorical features. Typically, these tables in industrial DLRMs contain trillions of parameters, necessitating model parallelism strategies to address memory constraints. However, as training systems expand with massive GPUs, the traditional fully parallelism strategies for embedding table post significant scalability challenges, including imbalance and straggler issues, intensive lookup communication, and heavy embedding activation memory. To overcome these limitations, we propose a novel two-dimensional sparse parallelism approach. Rather than fully sharding tables across all GPUs, our solution introduces data parallelism on top of model parallelism. This enables efficient all-to-all communication and reduces peak memory consumption. Additionally, we have developed the momentum-scaled row-wise AdaGrad algorithm to mitigate performance losses associated with the shift in training paradigms. Our extensive experiments demonstrate that the proposed approach significantly enhances training efficiency while maintaining model performance parity. It achieves nearly linear training speed scaling up to 4K GPUs, setting a new state-of-the-art benchmark for recommendation model training.