Source-linked AI summary
TensorDIMM: A Practical Near-Memory Processing Architecture for Embeddings and Tensor Operations in Deep Learning
Youngeun Kwon, Yunjae Lee, Minsoo Rhu
TL;DR
Embedding layers in recommender systems can exceed accelerator memory capacity and are constrained by memory bandwidth, making lookup and tensor operations costly. The paper combines near-memory processing on commodity-DIMM-based TensorDIMMs with a GPU-interfaced disaggregated TensorNode memory pool. Its prototype reports average 6.2−15.0× and 8.9−17.6× inference-time improvements over CPU-only and hybrid CPU-GPU implementations, respectively.
Problem
Embedding layers can reach several hundreds of GBs even for inference, exceeding GPU memory capacity and creating memory-capacity and bandwidth challenges.
Method
The paper proposes TensorDIMM, commodity-DIMM-based near-memory processing for embedding gathers and reductions, plus TensorNode, a GPU-interfaced disaggregated memory pool.
Results
6.2−15.0× and 8.9−17.6× average inference-time improvements are reported over CPU-only and hybrid CPU-GPU recommender-system implementations, respectively.
Takeaways & Limitations
TensorDIMM provides scalable memory capacity and bandwidth for embedding gathers and tensor operations in recommender systems.
Takeaways & Limitations
Because evaluation uses emulated TensorDIMMs, the paper cannot demonstrate model-quality improvements from training algorithms with larger embeddings.
Abstract
from arXiv · showhide
Recent studies from several hyperscalars pinpoint to embedding layers as the most memory-intensive deep learning (DL) algorithm being deployed in today's datacenters. This paper addresses the memory capacity and bandwidth challenges of embedding layers and the associated tensor operations. We present our vertically integrated hardware/software co-design, which includes a custom DIMM module enhanced with near-data processing cores tailored for DL tensor operations. These custom DIMMs are populated inside a GPU-centric system interconnect as a remote memory pool, allowing GPUs to utilize for scalable memory bandwidth and capacity expansion. A prototype implementation of our proposal on real DL systems shows an average 6.2-17.6x performance improvement on state-of-the-art recommender systems.
1 INTRODUCTION
Embedding layers create a DL memory wall because they require high capacity and bandwidth, while conventional architectures have focused mainly on dense computation. TensorDIMM addresses this with near-memory processing, custom tensor instructions, and a GPU-interfaced disaggregated memory pool.
- Embedding lookups and tensor manipulations are among the most memory-intensive DL workloads, reaching several hundreds of GBs even for inference.
- The paper targets embedding-layer capacity and bandwidth constraints in recommender systems through a vertically integrated hardware/software co-design.
- TensorDIMM enhances commodity buffered DIMMs with near-memory processing units for embedding gathers and reductions, reducing vector-fetch and reduction latency.
- TensorISA and its runtime provide bandwidth-efficient tensor operations, embedding address mapping, and scalable TensorDIMM utilization.
- TensorNode aggregates TensorDIMMs into a disaggregated memory pool connected inside a GPU-side, NVLINK-compatible high-bandwidth interconnect.
- 6.2−15.0× and 8.9−17.6× average inference-time improvements are reported over CPU-only and hybrid CPU-GPU recommender-system implementations, respectively.
2 BACKGROUND
DNN recommender systems combine sparse embedding features with dense features, making embedding lookup and tensor manipulation central operations. Their lookup tables scale with users and items, producing model footprints that can exceed GPU memory capacity.
- Commodity DRAM ranks and DIMMs provide the conventional modular organization for balancing memory capacity and bandwidth.
- DNN recommender systems use embeddings in applications including advertisements, social networking, recommendations, and news feeds.
- Sparse one-hot indices query embedding tables to produce dense vectors that are combined with other features before DNN processing.
- Embedding layers gather vectors from lookup tables and apply tensor concatenation or reductions such as additions, multiplications, and averages.
- Each user or item requires a unique embedding vector, so total embedding storage scales with the number of users and items.
3 MOTIVATION
Embedding lookup tables can reach several hundreds of GBs, exceeding GPU or NPU memory capacity and forcing systems to use bandwidth-limited CPU memory. This creates a trade-off between CPU computation and CPU-to-GPU communication, motivating scalable capacity and bandwidth expansion.
- Bandwidth-optimized stacked memory offers high bandwidth but only several tens of GBs, while embedding models require several hundreds of GBs.
- Embedding lookup tables are therefore stored in CPU memory, and embeddings are read across the CPU-to-GPU path for inference.
- Embedding gathers are memory-bandwidth-limited, so CPU-memory placement adds significant latency relative to an ideal GPU-only system.
- Larger embedding dimensions increase model size more dramatically than larger MLP dimensions in the evaluated recommender-system scaling study.
- CPU-only execution avoids PCIe transfer overhead but suffers from lower computation throughput, whereas hybrid CPU-GPU execution adds CPU-to-GPU communication latency.
- Future workloads are expected to use more lookups and larger embeddings, requiring scalable memory capacity and bandwidth expansion.
4 TENSORDIMM: AN NMP DIMM DESIGN FOR EMBEDDINGS & TENSOR OPS
TensorDIMM combines near-memory processing, scalable DIMM-level bandwidth, and TensorNode disaggregated memory to address embedding and tensor-operation bottlenecks. Its hardware/software design performs gathers and reductions near memory before transferring reduced tensors to GPUs over NVLINK.
- TensorDIMM architecture: Near-memory gathers and reductions reduce embedding-transfer volume by a factor of N before GPU processing.The design gathers and reduces N embeddings locally, then copies a single reduced tensor to GPU memory.
- TensorDIMM architecture: TensorDIMM adds NMP cores to buffered DIMMs for embedding gathers and tensor reductions while retaining commodity DRAM devices.The NMP design includes a DDR interface, vector ALU, and local SRAM queues; TensorDIMM can also operate as a normal buffered DIMM.
- Software architecture: TensorDIMM adds TensorISA primitives for GATHER, REDUCE, and AVERAGE, supported by a runtime that parallelizes operations across modules.The instruction and runtime layers encode embedding and tensor-operation information for near-memory execution.
- Scalable memory bandwidth: Rank-level address mapping distributes each embedding across TensorDIMMs so NMP bandwidth scales with the number of employed modules.The mapping enables NMP cores to work concurrently on distinct embedding subsets and uses the available DIMM-level parallelism.
- TensorNode system: TensorNode aggregates multiple TensorDIMMs into a disaggregated memory pool connected to GPUs through NVLINK and NVSwitch.This expands memory capacity independently of GPU-local memory and supports GPU reads and writes over the GPU-side interconnect.
- TensorNode system: NVLINK provides approximately 9× the communication bandwidth of PCIe, reducing TensorNode-to-GPU transfer latency.The high-bandwidth link complements local gather and reduction by accelerating movement of the reduced tensor between TensorNode and GPU.
5 EVALUATION METHODOLOGY
The evaluation combines cycle-level DRAM modeling with a proof-of-concept emulation on an NVIDIA DGX system. It assesses bandwidth utilization, system-level effects, and recommender-system performance across four embedding-based workloads under stated default configurations.
- Cycle-level simulation: The evaluation uses memory traces from DL frameworks and feeds GATHER, REDUCE, and AVERAGE transactions into the cycle-accurate Ramulator simulator.The simulation focuses on DRAM bandwidth utilization because TensorDIMM and TensorNode performance is bounded by memory-bandwidth use.
- Proof-of-concept prototype: The prototype emulates TensorNode with one DGX V100 GPU acting as TensorNode and another acting as a normal GPU.The setup uses DGX GPUs and NVLINK communication to approximate system-level effects of the proposed GPU-to-TensorNode connection.
- Benchmarks: The benchmark suite contains NCF, YouTube, Fox, and Facebook recommender systems using embeddings.These four neural-network-based applications are evaluated as representative recommender workloads.
- Benchmarks: The default evaluation uses batch size 64 and embedding dimension 512, with batch sizes from 1 to 128 in sensitivity studies.The default batch size is selected from a reported datacenter deployment range of 1–100.
- Area and power: Area and power overheads are evaluated using synthesized Verilog implementations and Micron’s DDR4 power calculator.The implementation targets a Xilinx Virtex UltraScale+ VCU1525 acceleration development board.
6 EVALUATION
The evaluation compares TensorDIMM-based TensorNode with CPU, hybrid, pooled-memory, and oracle GPU design points across bandwidth, latency, scaling, and interconnect conditions. TensorDIMM improves memory throughput and recommender-system performance while remaining relatively robust when communication bandwidth is reduced.
- Five recommender-system design points include CPU-only, CPU-GPU, PMEM, TensorDIMM-based TensorNode, and an oracle GPU-only system.
- 6.1 Memory Bandwidth Utilization: 4× average memory-bandwidth utilization improvement is achieved over the baseline CPU system, with maximum throughput of 808 versus 192 GB/sec.
- 6.1 Memory Bandwidth Utilization: Embedding dimension increases of up to 4× require proportionally more DIMMs, while the baseline CPU memory system saturates as capacity expands.
- 6.2 System-level Performance: TensorDIMM reduces both embedding lookup and embedding-copy latency in batch-64 inference.
- 6.2 System-level Performance: 84% average oracle-GPU performance is achieved, with at least 75% across design points and average speedups of 6.2× over CPU-only and 8.9× over CPU-GPU.
- 6.3 TensorDIMM with Large Embeddings: For larger embeddings, TensorNode achieves 6.2−15.0× and 8.9−17.6× improvements over CPU-only and CPU-GPU, respectively, with a maximum of 35×.
- 6.4 TensorDIMM with Low-bandwidth System Interconnects: With a 6× lower communication bandwidth, TDIMM experiences up to 15% performance loss, compared with PMEM’s maximum 68% loss.
7 RELATED WORK
The related-work discussion situates TensorDIMM among disaggregated memory, GPU-interconnect, DIMM-buffer acceleration, and DL accelerator research. It distinguishes TensorDIMM’s embedding-focused inference architecture from prior systems focused on other workloads or unable to fully exploit DRAM bandwidth for embeddings.
- Disaggregated memory commonly uses a remote pool connected over PCIe to increase CPU-accessible memory capacity.
- Prior GPU-interconnect disaggregation work focuses on DL training, whereas TensorDIMM primarily targets DL inference.
- Prior DIMM-buffer accelerators add custom logic, but the cited MCN approach may be suboptimal for embedding tensor operations because it cannot maximally use DRAM bandwidth.
- Figures 14–16 cover recommender-system performance, larger embeddings, and communication-bandwidth sensitivity, respectively, averaged or normalized as specified in their captions.
- The paper contrasts its embedding-focused architecture with broader single-GPU/NPU DL accelerators and scale-out or network-centric DL training platforms.
8 CONCLUSION
TensorDIMM combines near-memory processing with commodity DRAM in a disaggregated memory pool to address embedding-layer memory capacity and bandwidth constraints. It reports substantial recommender-system speedups over CPU-only and hybrid CPU-GPU implementations.
- TensorDIMM combines NMP cores with commodity DRAM devices to accelerate DL tensor operations.
- 6.2−15.0× and 8.9−17.6× average performance improvements are achieved over conventional CPU-only and hybrid CPU-GPU recommender-system implementations, respectively.
- Built on a disaggregated memory pool, TensorDIMM provides memory-capacity and bandwidth scaling for embeddings.
- The paper presents TensorDIMM as the first quantitative exploration of architectural solutions tailored for embeddings and tensor operations.