TRACI: Network Acceleration of Input-Dynamic Communication for Large-Scale Deep Learning Recommendation Model

Large-scale deep learning recommendation models (DLRMs) rely on embedding layers with terabyte-scale embedding tables, which present significant challenges to memory capacity. In addition, these embedding layers exhibit sparse and random data access patterns, which demand high memory bandwidth. Multi-GPU systems provide a promising solution, allowing for the scaling of both memory and aggregated bandwidth. However, network communication bandwidth becomes a bottleneck for multi-GPU DLRM systems. Overcoming the communication bottleneck is crucial to unlocking the potential of multi-GPU systems for efficient and high-performance DLRM training.

This paper introduces TRACI, an in-network acceleration architecture designed to optimize the communication operator in embedding layers: Aggregation. While in-network acceleration has proven successful for the All-Reduce communication collective, existing solutions do not directly apply to Aggregation due to two key challenges. Firstly, existing multi-GPU shared memory operations are designed for point-to-point communication and do not allow the network to proactively optimize communication. Secondly, in Aggregation, data transfer patterns are dynamic and dependent on input, demanding the network to dynamically discover and exploit message connections on-the-fly. To address these challenges, we propose a solution that involves a novel network transaction and switch hardware design. We introduce a new network transaction that augments messages with input reuse and output reuse identifications, and can empower the network to proactively reduce