GeoTrain improves model training speed and energy use on mixed GPUs
GeoTrain: Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources
Distributed, Parallel, and Cluster Computing
Summary
Training large AI models usually needs expensive, super-fast computer clusters in one place. The authors look at training on many slower, different GPUs spread worldwide, which is cheaper but harder to manage. They built GeoTrain, a system that picks the best GPUs, splits the model smartly, and plans work to save time and electricity. Tests show GeoTrain can speed up training by over 60% and save nearly 12% energy compared to older methods. It also finds the best plans quickly and scales well with more GPUs.
What this means in practice
- •For cloud infrastructure teams: Configure training across geographically spread GPUs balancing speed and energy for large AI models.
- •For energy-conscious ai developers: Build training workflows that reduce power consumption while maintaining or improving model training time.
Authors
Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi
Abstract
Training foundation models at the frontier typically requires large-scale dedicated and centralized clusters with fast interconnects - infrastructure that is prohibitively expensive and generally accessible only to well-resourced organizations. Decentralized training over pooled, geographically distributed GPUs offers an alternative, but introduces substantial heterogeneity in compute, memory, power, and network performance. Existing distributed training systems are primarily designed for clusters with fast, low-latency interconnects and typically use all available GPUs when configuring parallel training. In geo-distributed environments, however, slower devices and weak network links can become bottlenecks, while optimizing solely for throughput overlooks significant energy cost. We present GeoTrain, an automatic parallelism system that jointly optimizes training performance and energy consumption for heterogeneous geo-distributed GPU environments. GeoTrain simultaneously selects GPUs, partitions model layers, and determines pipeline configurations while accounting for heterogeneous compute, memory, power, and network constraints. To efficiently explore this large combinatorial search space, GeoTrain employs an optimized memetic algorithm with specialized genetic operators and local search. Experiments with decoder-only Transformer models on real-world and simulated heterogeneous topologies demonstrate that GeoTrain reduces iteration time and energy consumption by up to $61.4\%$ and $11.7\%$, respectively, over the existing state-of-the-art approaches. It also finds the exact optimum on verifiable instances and scales beyond exact methods as the number of GPUs increases, while exposing favorable energy-time trade-offs, including cases where its time-optimal distributed training plan is both faster and more energy-efficient than the baselines.