Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ProactivePIM: Accelerating Weight-Sharing Embedding Layer with PIM for Scalable Recommendation System

Youngsuk Kim, Junghwan Lim, Hyuk-Jae Lee, Chae Eun Rhee

TL;DR

The paper tackles the embedding-layer bottleneck in large-scale DL-based recommender systems, driven by growing model sizes and memory bandwidth demands. It introduces ProactivePIM, a two-level PIM/HBM-DIMM architecture that exploits intra-GnR locality and subtable size imbalance to accelerate weight-sharing embedding algorithms (QR-trick and TT-Rec) while eliminating CPU-PIM communication via subtable duplication and SRAM caching. Key contributions include a table-wise prefetch scheme, a subtable mapping strategy, and a TT-Rec optimization that converts CnR into two-stage skinny GEMMs, achieving up to 2.89x/2.78x speedups over baselines and surpassing prior NMP designs by substantial margins. The approach enables scalable, memory-bound inference for large embedding tables and can extend to other memory-intensive workloads by leveraging static caching and the two-level PIM design.

Abstract

Although deep learning-based personalized recommendation systems provide qualified recommendations, they strain data center resources. The main bottleneck is the embedding layer, which is highly memory-intensive due to its sparse, irregular access patterns to embeddings. Recent near-memory processing (NMP) and processing-in-memory (PIM) architectures have addressed these issues by exploiting parallelism within memory. However, as model sizes increase year by year and can exceed server capacity, inference on single-node servers becomes challenging, necessitating the integration of model compression. Various algorithms have been proposed for model size reduction, but they come at the cost of increased memory access and CPU-PIM communication. We present ProactivePIM, a PIM system tailored for weight-sharing algorithms, a family of compression methods that decompose an embedding table into compact subtables, such as QR-trick and TT-Rec. Our analysis shows that embedding layer execution with weight-sharing algorithms increases memory access and incurs CPU-PIM communication. We also find that these algorithms exhibit unique data locality characteristics, which we name intra-GnR locality. ProactivePIM accelerates weight-sharing algorithms by utilizing a heterogeneous HBM-DIMM memory architecture with integration of a two-level PIM system of base-die PIM (bd-PIM) and bank-group PIM (bg-PIM) inside the HBM. To gain further speedup, ProactivePIM prefetches embeddings with high intra-GnR locality into an SRAM cache within bg-PIM and eliminates the CPU-PIM communication through duplication of target subtables across bank groups. With additional optimization techniques, our design effectively accelerates weight-sharing algorithms, achieving 2.22x and 2.15x speedup in QR-trick and TT-Rec, respectively, compared to the baseline architecture.

ProactivePIM: Accelerating Weight-Sharing Embedding Layer with PIM for Scalable Recommendation System

TL;DR

The paper tackles the embedding-layer bottleneck in large-scale DL-based recommender systems, driven by growing model sizes and memory bandwidth demands. It introduces ProactivePIM, a two-level PIM/HBM-DIMM architecture that exploits intra-GnR locality and subtable size imbalance to accelerate weight-sharing embedding algorithms (QR-trick and TT-Rec) while eliminating CPU-PIM communication via subtable duplication and SRAM caching. Key contributions include a table-wise prefetch scheme, a subtable mapping strategy, and a TT-Rec optimization that converts CnR into two-stage skinny GEMMs, achieving up to 2.89x/2.78x speedups over baselines and surpassing prior NMP designs by substantial margins. The approach enables scalable, memory-bound inference for large embedding tables and can extend to other memory-intensive workloads by leveraging static caching and the two-level PIM design.

Abstract

Although deep learning-based personalized recommendation systems provide qualified recommendations, they strain data center resources. The main bottleneck is the embedding layer, which is highly memory-intensive due to its sparse, irregular access patterns to embeddings. Recent near-memory processing (NMP) and processing-in-memory (PIM) architectures have addressed these issues by exploiting parallelism within memory. However, as model sizes increase year by year and can exceed server capacity, inference on single-node servers becomes challenging, necessitating the integration of model compression. Various algorithms have been proposed for model size reduction, but they come at the cost of increased memory access and CPU-PIM communication. We present ProactivePIM, a PIM system tailored for weight-sharing algorithms, a family of compression methods that decompose an embedding table into compact subtables, such as QR-trick and TT-Rec. Our analysis shows that embedding layer execution with weight-sharing algorithms increases memory access and incurs CPU-PIM communication. We also find that these algorithms exhibit unique data locality characteristics, which we name intra-GnR locality. ProactivePIM accelerates weight-sharing algorithms by utilizing a heterogeneous HBM-DIMM memory architecture with integration of a two-level PIM system of base-die PIM (bd-PIM) and bank-group PIM (bg-PIM) inside the HBM. To gain further speedup, ProactivePIM prefetches embeddings with high intra-GnR locality into an SRAM cache within bg-PIM and eliminates the CPU-PIM communication through duplication of target subtables across bank groups. With additional optimization techniques, our design effectively accelerates weight-sharing algorithms, achieving 2.22x and 2.15x speedup in QR-trick and TT-Rec, respectively, compared to the baseline architecture.
Paper Structure (25 sections, 1 equation, 13 figures, 1 table)

This paper contains 25 sections, 1 equation, 13 figures, 1 table.

Table of Contents

  1. Introduction
  2. Background
  3. Personalized Recommendation System
  4. NMP/PIM for Recommendation System
  5. Weight-sharing Algorithms for Embedding Table Reduction
  6. Motivation
  7. Weight-Sharing Algorithm Analysis
  8. Chances and Challenges of PIM Implementation.
  9. Opportunities for Enhancing PIM Speedup
  10. ProactivePIM Architecture
  11. ProactivePIM Architecture
  12. Hot Subembedding Prefetch
  13. Subtable Mapping Strategy
  14. TT-Rec Optimization on ProactivePIM
  15. ProactivePIM Implementation
  16. ...and 10 more sections

Figures (13)

  • Figure 1: Model size and bandwidth growth. (Reproduced from MTIA firoozshahian2023mtia and RecShard sethi2022recshard)
  • Figure 2: DLRM architecture
  • Figure 3: Weight-sharing algorithms. (a) QR-trick with hash collision set to 4 (b) TT-Rec with total TT cores set to 3.
  • Figure 4: Rooline analysis and CPU-PIM communication overview. (a) Roofline analysis. (Evaluated on NVIDIA Tesla V100 with varying batch sizes and total embeddings per GnR.) (b) CPU-PIM communication example (c) Execution time breakdown on weight-sharing algorithms when bankgroup level PIM is utilized. Ideal throughput stands for throughput without CPU-PIM communication.
  • Figure 5: ProactivePIM system overview. (a) Host side implementation and heterogeneous memory architecture. MC indicates memory controller. (b) 2-level PIM system overview within HBM2.
  • ...and 8 more figures