STEMS: Spatio-Temporal Scheduling to Accelerate LLM Training in Optically Interconnected Clusters

Abstract

All-optical interconnects (AOIs) provide large yet reconfigurable bandwidth with high energy efficiency, making them key enablers for sustaining the rapid expansion of large language model (LLM) training clusters. However, the flexibility of AOIs can never be fully exploited if not jointly optimized with the collective communications of LLM training jobs, which often exhibit high skewness in both time and spatial domains. To address this challenge, we propose STEMS, a SpatioTEMporal Scheduling system for running LLM training jobs in an optically-interconnected AI cluster. Specifically, STEMS optimizes both spatial (collective communication topologies, topology engineering of AOI and traffic engineering) and temporal (job grouping for communication interleaving) scheduling to effectively accelerate LLM training jobs. We formulate a mixed integer linear programming model to precisely model the optimization in STEMS and develop polynomial-time heuristics for practical operations of STEMS. Being prototyped over an optically-interconnected GPU cluster running real LLM training jobs, STEMS is shown to reduce training iteration time by up to 34.32% over the state-of-the-art, with improved GPU utilization. Extensive large-scale simulations (up to 128 racks and $1,000$ concurrent jobs) further confirm the scalability and superiority of STEMS.