In an effort to enhance the efficiency of large-scale recommendation systems, NVIDIA has introduced EMBark, a novel approach aimed at optimizing embedding processes in deep learning recommendation models. According to NVIDIA, recommendation systems are pivotal to the Internet industry, and efficiently training them poses a significant challenge for many companies.

Challenges in Training Recommendation Systems

Deep learning recommendation models (DLRMs) often incorporate billions of ID features, necessitating robust training solutions. Recent advancements in GPU technology, such as NVIDIA Merlin HugeCTR and TorchRec, have improved DLRM training by utilizing GPU memory to handle large-scale ID feature embeddings. However, with an increase in the number of GPUs, the communication overhead during embedding has become a bottleneck, sometimes accounting for over half of the total training overhead.

EMBark’s Innovative Approach

Presented at RecSys 2024, EMBark addresses these challenges by implementing 3D flexible sharding strategies and communication compression techniques, aiming to balance the load during training and reduce communication time for embeddings. The EMBark system includes three core components: embedding clusters, a flexible 3D sharding scheme, and a sharding planner.

Embedding Clusters

These clusters group similar features and apply customized compression strategies, facilitating efficient training. EMBark categorizes clusters into data parallel (DP), reduction-based (RB), and unique-based (UB) types, each suited for different training scenarios.

Flexible 3D Sharding Scheme

This innovative scheme allows for precise control of workload balance across GPUs, utilizing a 3D tuple to represent each shard. This flexibility addresses the imbalance issues found in traditional sharding methods.

Sharding Planner

The sharding planner employs a greedy search algorithm to determine the optimal sharding strategy, enhancing the training process based on hardware and embedding configurations.

Performance and Evaluation

EMBark’s efficacy was tested on NVIDIA DGX H100 nodes, demonstrating significant improvements in training throughput. Across various DLRM models, EMBark achieved an average 1.5x increase in training speed, with some configurations reaching up to 1.77x faster than traditional methods.

By enhancing the embedding process, EMBark significantly improves the efficiency of large-scale recommendation system models, setting a new standard for deep learning recommendation systems. For more detailed insights into EMBark’s performance, you can view the research paper.

Image source: Shutterstock