Efficient and Automated Scaling of Distributed AI Infrastructure
COM2 Level 4
Executive Classroom, COM2-04-02
Abstract:
Recent advances in AI have been driven by model scaling. While scaling leads to predictable improvements, it also introduces new infrastructure challenges. In the pre-training phase, carefully orchestrated model parallelism decisions must be made to maximize training throughput; while in the post-training and inference phase, serving highly variable requests with dynamic length create memory pressure. This thesis develops automated methods to scale training and inference efficiently on distributed systems by addressing three such challenges.
First, distributed data parallel training workloads typically demands stable and high-bandwidth network. However, such networks are often unavailable in commercial cloud environments, which suffers from network jittering and congestion. To mitigate this, we propose a model-system co-design approach that leverages asynchronous training and congestion-aware data pipeline to accelerate generative adversarial network training. Demonstrating this approach, we present ParaGAN, a cloud-based training system achieving 91% scaling efficiency across 1024 accelerators.
Second, as models scale in size, existing approaches to identify model parallel strategies at scale are prohibitively expensive.Observing that neural network architectures often contain inherent redundancies, we introduce TAPAS, an automated tensor parallelization framework. TAPAS dramatically accelerates the search process by two orders of magnitude, while achieving competitive performance compared with state-of-the-art expert-engineered libraries.
Last but not least, for reasoning and agentic workloads, we identify a cascading preemption pattern: under GPU memory pressure, uninformed victim selection triggers back-to-back preemptions, wasting compute and harming latency. We mitigate this with (i) an online predictor that guides memory-aware victim selection and (ii) a fine-grained, asynchronous KV-cache offload mechanism that overlaps transfers with GPU computation. Results show that our method can optimize both preemption frequency and tail latency.
Collectively, our studies advance automated and efficient scaling of distributed training and inference on real-world infrastructure, resulting in production features now available on public cloud platforms.