Arachne plans efficient machine learning training on mixed GPU clusters
Arachne: Learning to Plan Parallel Training on Dynamic Heterogeneous Clusters
Distributed, Parallel, and Cluster Computing
Summary
Training big AI models on systems with many different GPUs is hard because the available hardware changes and the GPUs are not all the same. The authors created Arachne, a planner that learns how to organize training tasks to run well even when the hardware is mixed and changing. It focuses on finding good structures for the training process, making decisions faster and better. Testing shows Arachne often performs as well or better than other methods, sometimes boosting training speed by a lot.
What this means in practice
- •For cloud platform engineers: Optimize training schedules for large AI models on clusters with varied GPU types to increase hardware utilization and throughput.
- •For data center operators: Adapt training task planning dynamically as GPU availability changes to maintain efficient AI training.
Authors
Taeyoon Kim, Yonguk Song, Seoyeong Choy, Hexiao Duan, Dong Li, Seo Jin Park, Myeongjae Jeon
Abstract
Training large machine learning models on shared GPU infrastructures faces two challenges: (1) GPU availability shifts dynamically with varying resource demands from tenants, and (2) hardware heterogeneity accumulates as datacenters continuously adopt new GPU generations. Due to the vast search space induced by heterogeneous GPU types and node sizes, training planners must prune it aggressively to remain tractable, yet must also derive high-throughput plans promptly as cluster configurations change. Arachne achieves this goal through a learning-based planner that reduces the full planning problem to a search over pipeline structures. Arachne encapsulates planning decisions in a structural template and learns to construct plans from templates over diverse cluster configurations offline. This design is effective because structural decisions constitute the performance-critical core of a parallelism plan, while the rest follows by rule or from a small priced candidate set once the plan structure is fixed. Evaluation shows that Arachne matches or exceeds the best plan found by five existing planners across clusters with varying GPU types and node sizes for three models of different sizes by up to 84.5% in throughput on dense models and 4.6x on MoE models.