Shesha: Opportunistic In-network Acceleration of Asynchronous Distributed Reinforcement Learning

Abstract

Large-scale training for distributed Machine Learning can cause congestion at bottleneck switch ports, leading to model staleness through update losses. This is particularly detrimental for asynchronous Distributed Reinforcement Learning (DRL) training, as stale updates are known to degrade convergence performance in asynchronous settings. This paper presents \textit{Shesha}, an in-network DRL accelerator engine, which \textit{opportunistically aggregates} asynchronously generated model updates \textit{on the fly} while they traverse the data plane queue. This aggregation operation motivates an alternative queue design, which we prototype and envision for future Top-of-Rack switches. We further present corresponding host-side transmission control in the face of possible congestion, taking advantage of in-network accelerator feedback. A quantification of model staleness, denoted Age-of-Model (AoM), together with a formal verifier allows us to reason on system-wide AoM objectives in multi DRL-cluster scenarios. Shesha shows significant reductions in model staleness and queue congestion, improving overall convergence behavior for asynchronous DRL workloads.