GeoTrain: Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources

GeoTrain: Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi∗ Faculty of Information Technology, Monash University, Melbourne, Australia

arXiv:2610.26178v1 [cs.DC] 19 Oct 2026

Abstract

Computing methodologies → Parallel computing methodologies; Machine learning; Parallel computing methodologies; Machine learning.

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. Efficiently selecting and configuring a subset of heterogeneous GPUs therefore requires jointly reasoning about resource characteristics, model partitioning, and pipeline placement. 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.

Keywords: decentralized training, distributed training, foundation models, pipeline parallelism, model partitioning, heterogeneous systems, communication-aware optimization, memetic algorithm

1

Training foundation models, such as GPT-3 and LLaMA, requires distributing computation across hundreds or thousands of GPUs [2, 29]. As models continue to scale, organisations are increasingly pooling all available GPU resources, including heterogeneous hardware spanning multiple generations and geographically distributed clusters linked by wide-area networks with diverse bandwidth and latency characteristics. Making efficient use of such infrastructure is essential to reducing both training time and the considerable energy expenditure of large-scale model training. Systems such as Alpa [35] and AMP [9] have advanced the automation of parallelism planning for distributed training. Also, a separate line of work, DTFM [34], SWARM [22], Sailor [25], Metis [30], and CocktailSGD [32], has begun to address training over heterogeneous or geo-distributed resources directly. However, across these efforts, three major assumptions persist that limit their applicability to the heterogeneous and geo-distributed paradigm: Full-GPUs participation. Existing systems assume fullcluster participation, where all available GPUs are always used for training. This reduces the problem to optimizing configurations over a fixed set of devices, avoiding the combinatorial complexity of GPU subset selection. However, this assumption becomes inefficient in heterogeneous, geodistributed settings: slow GPUs and weak interconnects can dominate iteration time and negate the benefit of additional parallelism. As shown in Fig. 2a, selecting an optimal subset of 𝑘 = 4 GPUs from 7 heterogeneous devices reduces iteration time by 75% compared to using the full cluster. This demonstrates that carefully chosen subsets can outperform full-cluster execution in heterogeneous environments. Uniform stage partitioning. Prior systems partition model layers uniformly across pipeline stages, implicitly assuming homogeneous compute capacity across GPUs. This assumption breaks under heterogeneity, where uniform assignment

CCS Concepts: • Computer systems organization → Distributed architectures; Distributed architectures; •

∗ Corresponding author.

Emails: khanh.hoang@monash.edu, mohammad.goudarzi@monash.edu.

Introduction

tianzhang.he@monash.edu, 1

Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi

7000 6000 5000 4000 3000 2000 1000

Optimal subset

All GPUs

Iteration time (ms)

Iteration time (ms)

Figure 1. A high-level comparison between GeoTrain and a standard uniform planer using all available GPUs.

optimal k=4

1

2

3

4

5

Subset size k

6

7

1200 1053 ms 1000 800 600 400 200 0 Uniform

candidate by both iteration time and energy under the heterogeneous, geo-distributed setting, so the entire search is driven toward configurations that jointly minimize iteration time and energy. The MA evolves chromosomes using specialized genetic operators (i.e., crossover, mutation, and local search). The crossover cuts at training stage boundaries so offspring inherit whole stages from each parent, greedily keeps the fittest cut, and occasionally allows a random mid-stage cut for diversity. The mutations employ different strategies at varying granularity to refine solutions gradually, escape local optima, and broaden exploration. A local search then finds the training pipeline’s bottleneck and rebalances the workload to relieve it. Figure 1 illustrates the resulting scheduling effect by comparing GeoTrain with standard uniform planning using all available GPUs. Together, these components let GeoTrain explore the plan space efficiently and converge quickly toward balanced, low-cost pipelines. In designing GeoTrain, we make following contributions: • Heterogeneity-aware Pipeline Planning Unlike prior systems that typically assume all available GPUs must be used, GeoTrain jointly selects a subset of GPUs and determines a non-uniform layer partition across pipeline stages. It accounts for heterogeneous compute capacities and asymmetric network links to avoid slow devices and poorly connected links, while balancing per-stage cost to eliminate stragglers. This enables GeoTrain to construct balanced pipelines that achieve lower iteration time than full-cluster or uniform partitioning. • Time-Energy-aware optimization. We formulate parallelism planning as an optimization that jointly minimizes iteration time and energy consumption. An analytical energy cost model captures GPU-specific power profiles under the 1F1B pipeline schedule [14, 15], while a user-tunable weighting parameter allows operators to navigate the time-energy trade-off. By varying the parameter, GeoTrain can produce energy-saving, balanced, or time-optimal training configurations, enabling operators to select operating points along the Pareto frontier according to their performance and energy requirements.

526 ms

Non-Uniform

(a) Selecting an optimal GPU sub- (b) Non-uniform layer assignset (𝑘=4) yields lower iteration ment eliminates the bottleneck time than using all GPUs. stage under uniform one.

Figure 2. Motivating observations on a real-world heterogeneous geo-distributed configuration with 7 GPUs

creates straggler stages on weaker or bandwidth-constrained devices that stall the entire pipeline. Capacity-aware nonuniform partitioning that matches layer allocation to GPU compute and memory constraints yields a more balanced pipeline and significantly lower iteration time (Fig. 2b). Energy-unaware optimization. Existing automatic parallelism systems optimize solely for training throughput or iteration time, ignoring energy consumption. However, as the environmental and financial costs of large-scale training continue to rise [20, 26], energy efficiency has become a critical objective. Importantly, equivalent training speeds can exhibit substantially different energy profiles across pipeline configurations due to variations in GPU selection, layer placement, and micro-batch sizing, especially in heterogeneous environments with diverse power characteristics. In this paper, we present GeoTrain, a system that automatically discovers efficient pipeline parallelism plans for training on heterogeneous, geo-distributed GPUs. To address limitations of prior work, it employs an optimized Memetic Algorithm (MA) with domain-specific operators to co-optimize three coupled decisions: which GPUs to use, how to partition model layers, and where to place each partition. GeoTrain encodes a plan as an integer chromosome, capturing all three decisions in one encoding. A key design is its fitness function: an analytical model that scores each 2

GeoTrain : Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources

Table 1. Comparison of GeoTrain with heterogeneous and geo-distributed foundation-model training systems.

• Domain-specific Memetic Algorithm. We design a domain-specific Memetic Algorithm tailored to the combinatorial planning space of heterogeneous, geo-distributed pipeline training. It combines stage-aware crossover, multigranularity mutation, and bottleneck-driven local search to efficiently explore candidate configurations and refine promising solutions. • Extensive real-world experiments. We evaluate GeoTrain using both real-world experiments and simulation on heterogeneous geo-distributed GPU configurations with Transformer models examining scalability, optimality, network sensitivity, and energy-time trade-offs. The remainder of this paper is organized as follows. We analyze distributed model training (Sect. 2) and related work (Sect. 3). Section 4 formalizes the problem. Section 5 presents GeoTrain design. Section 6 analyzes the system and the results. Section 7 concludes and discusses future work.

2

Work

Environment

Search Space

Comp.

Net.

Asym.

Non-unif.

GPU

hetero.

aware

links

partition

subset

[35]

✗

✓

✗

✓

✗

[9]

✓

✓

✗

✓

✗

[19]

✓

✗

✗

✗

[22]

✓

✓

✗

[34]

✗

✓

[36]

✓

✓

[25]

✓

✓

[30]

✓

Ours

✓

Method

Objective T

E

ILP+DP

✓

✗

DP

✓

✗

✗

Manual

✓

✗

✗

✗

Dynamic

✓

✗

✗

✗

✗

EA

✓

✗

✗

✓

✗

MILP+H

✓

✗

✗

✓

✓

DP+H

✓

✗

✗

✗

✓

✗

DFS

✓

✗

✓

✓

✓

✓

MA

✓

✓

T : iteration time; E: energy. ILP: integer linear programming; MILP: mixedinteger linear programming; DP: dynamic programming; EA: evolutionary algorithm; DFS: depth-first search; H : heuristic pruning; MA: memetic algorithm. GPU subset: fewer than all available GPUs may be selected.

Preliminaries

Training foundation models involves repeatedly processing training batches through forward and backward passes to compute gradients. When models exceed single-GPU memory, distributed training partitions computation across multiple GPUs using parallelism strategies.

warmup and drain, while the second reflects the steady-state throughput determined by the bottleneck stage. 2.3

We consider a decoder-only Transformer [31], a widely used foundation-model architecture, with hidden size 𝐻 and 𝐿 layers. Each layer consists of a multi-head self-attention block with 4 𝐻 × 𝐻 projections followed by a feed-forward network (FFN) with structure 𝐻 → 4𝐻 → 𝐻 . The parameter count per layer is:

2.1 Parallelism Methods Distributed training commonly uses three forms of parallelism: Pipeline parallelism (PP) partitions model layers into stages, each mapped to different devices [4, 14]. Microbatches flow through the pipeline, with communication occurring at stage boundaries for activations and gradients. We focus on PP because its communication is confined to stage boundaries, making it well suited to geo-distributed environments compared with Data parallelism (DP) [10], which required frequent gradient synchronization, and Tensor parallelism (TP) [24], which incurs substantial fine-grained intralayer communication overhead. The work could be extended to incorporate DP and TP to exploit high-speed interconnects between co-located GPUs, but this is outside our current focus and left for future work.

𝑃ℓ ≈ 4𝐻 2 + 8𝐻 2 = 12𝐻 2,

𝑄 ∑︁

𝜏𝑞 (𝑢) + (𝑀 − 1) · max 𝜏𝑞 (𝑢)

(2)

where the first term covers the attention projections and the second covers the FFN. Biases and LayerNorm parameters are 𝑂 (𝐻 ) and omitted. The matrix multiplication (matmul) cost to process one token through one layer in the forward pass is approximately [15]: 𝐶 ℓ ≈ 2𝑃ℓ FLOPs/token,

(3)

The backward pass requires roughly 2× the forward FLOPs [7].

3

2.2 Pipeline Schedules Given a pipeline of 𝑄 stages processing 𝑀 micro-batches, the pipeline schedule determines the ordering of forward and backward computations across stages. In this work, we adopt the 1F1B schedule [14, 15], which gives the following iteration time formula: 𝑇 iter (𝑀, 𝑢) =

Decoder-Only Transformer Architecture

Related Work

In this section, we review relevant work, identify its gaps, and position our contributions. Automatic parallelism planning. Alpa [35] automates inter- and intra-operator parallelism using dynamic programming and integer linear programming, respectively. AMP [9] uses a heterogeneity-aware cost model for uneven layers and heterogeneous accelerators, while Galvatron [11] automates hybrid parallelism via dynamic programming. These works automate parallelism planning but do not address geo-distributed topology or energy-aware optimization. Heterogeneous and geo-distributed training. DTFM [34] addresses heterogeneous networks with explicit bandwidth

(1)

𝑞=1,...,𝑄 𝑞=1

where 𝜏𝑞 (𝑢) = 𝜏𝑞𝐹 (𝑢)+𝜏𝑞𝐵 (𝑢) is the forward-backward service time of stage 𝑞 for micro-batch size 𝑢. The first term captures the latency of a micro-batch through the pipeline, including 3

Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi

capacity (GB) and 𝑈𝑖mem ∈ [0, 1] is the fraction usable after framework and runtime overhead. Idle and active power: 𝑃𝑖idle and 𝑃𝑖act (watts), where 𝑃𝑖act is the power drawn when computing at the effective compute capacity 𝐹˜𝑖 . For each directed edge (𝑔𝑖 , 𝑔 𝑗 ) ∈ 𝐸: a communication link from GPU 𝑖 to 𝑗, we define: Effective bandwidth: 𝐵˜ 𝑖 𝑗 = 𝛽𝑖 𝑗 · 𝐵𝑖 𝑗 , where 𝐵𝑖 𝑗 is the nominal link bandwidth (Gbps) and 𝛽𝑖 𝑗 ∈ [0, 1] is its utilization. One-way latency: 𝐿𝑖 𝑗 , the time required for data to travel from GPU 𝑖 to 𝑗 over the link (ms).

and latency, but operates on a given parallelization plan. SWARM [22] targets poorly connected, heterogeneous devices through randomized, adaptively rebalanced pipelines, while CocktailSGD [32] reduces communication through gradient compression. Sailor [25] and HGTrainer [36] address heterogeneous geo-distributed resources; HetPipe [19] combines pipelined model and data parallelism across heterogeneous GPUs. Oobleck [5] focuses on fault-tolerant pipeline re-instantiation. Metis [30] automatically partitions workloads across heterogeneous GPUs but does not target geodistributed topology or energy-aware planning. Pipeline parallelism and memory. GPipe [4] popularizes micro-batch pipelining, while PipeDream [14] introduces the 1F1B schedule adopted by many later works. GraphPipe [6] introduces graph pipeline parallelism for general DNN computation graphs by jointly optimizing graph partitioning and micro-batch scheduling. AdaPipe [27] jointly optimizes recomputation and stage partitioning for compute-memory balance, rather than network and hardware heterogeneity. ZeRO [21] and Megatron-LM [24] target intra-cluster parallelism, complementary to our WAN-pipeline focus. Energy-aware foundation model training. Zeus [33] studies GPU power and energy consumption, motivating our energy-aware planning model. Strubell et al. [26] and Patterson et al. [20] quantify training’s energy and carbon footprint. However, these works do not incorporate energy into automatic parallelism planning. Across these lines of work, no existing planner combines the capabilities required for heterogeneous, geo-distributed training, as Table 1 summarizes. GeoTrain addresses these gaps by supporting directed network topologies with asymmetric links, jointly selecting GPU subsets and non-uniform layer partitions, and incorporating energy into plan search.

4

4.2 Memory Model We consider a decoder-only Transformer. The training dataset is split into batches of 𝐵 samples each. Each batch is further divided into 𝑀 micro-batches of size 𝑢 = 𝐵/𝑀 samples, where each sample is a sequence of 𝑇 tokens. Per-layer memory. The training-time memory footprint of one Transformer layer consists of parameters with optimizer state plus activations stored for backpropagation: Memℓtrain ≈ 4 𝑃ℓ 𝜅 + (𝑄 + 1) 𝛼 𝜅 𝑢 𝑇 𝐻

where 𝜅 is the bytes per element (4 for FP32, 2 for FP16); 4𝑃ℓ 𝜅 covers the memory size of parameters, gradients, and the two Adam moment estimates [8]; 𝛼 ≈ 10 is the number of intermediate tensors retained per layer for backpropagation; and (𝑄 + 1) bounds the number of micro-batches whose activations the worst-case stage retains during 1F1B warmup. Per-stage memory. For a stage containing 𝑛𝑞 consecutive layers, the total memory footprint is the per-layer cost scaled by the number of layers:
Mem𝑞train ≈ 𝑛𝑞 4 𝑃ℓ 𝜅 + (𝑄 + 1) 𝛼 𝜅 𝑢 𝑇 𝐻 (5) Last-stage memory. The final stage incurs extra terms of language model head projection (𝐻 → 𝑉 ) and its output: train Memtrain + last ≈ Mem𝑞

System Model and Problem Formulation

This section formalizes GeoTrain’s planning problem. We model resource constraints and costs, define the solution space, and formulate the optimization problem. 4.1

(4)

4 𝑃 lm 𝜅 | {z } LM head + Adam

+ 2𝜅 𝑢𝑇 𝑉 | {z }

(6)

logits + gradient

where 𝑃lm = 𝐻 × 𝑉 is the LM head parameter count and 𝑉 is the vocabulary size. Eq. 5 and Eq. 6 determines the memory constraint enforced in Eq. 19.

Hardware Model 4.3

We model the GPU network as a directed graph G = (𝐺, 𝐸), where 𝐺 = {𝑔1, 𝑔2, . . . , 𝑔𝑁 } is the set of 𝑁 available GPUs and 𝐸 is the set of directed communication links between them. It captures three key characteristics of geo-distributed environments often overlooked in prior work: heterogeneous GPU performance, asymmetric WAN bandwidth, and nonuniform latencies ranging from intra-rack microseconds to intercontinental hundreds of milliseconds. Each GPU 𝑔𝑖 ∈ 𝐺 is characterized by three properties: comp Effective compute capacity: 𝐹˜𝑖 = 𝑈𝑖 · 𝐹𝑖 , where 𝐹𝑖 comp is the peak compute capacity (FLOPS) and 𝑈𝑖 ∈ [0, 1] is the fraction effectively available for training. Effective memory: 𝑀˜ 𝑖 = 𝑈𝑖mem · 𝑀𝑖 , where 𝑀𝑖 is the total memory

Performance Model

We derive the per-stage service time used in the 1F1B iterationtime formula (Eq. 1). Stage computation time. Stage 𝑞 with 𝑛𝑞 layers placed on GPU 𝑔𝑞 processes one micro-batch of size 𝑢 through forward and backward passes in time comp

𝜏𝑞

(𝑢) =

(1 + 𝛾) 𝑛𝑞 𝐶 ℓ 𝑢 𝑇 , 𝐹˜𝑔𝑞

(7)

where 𝛾 ≈ 2 accounts for the backward pass costing roughly twice the forward, 𝐶 ℓ ≈ 2𝑃ℓ ≈ 24𝐻 2 is the per-token forward compute cost of one Transformer layer, and 𝐹˜𝑔𝑞 is the effective compute capacity of the assigned GPU. The last stage 4

GeoTrain : Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources

adds the LM head projection: comp

𝜏𝑄

(𝑢) =

Total pipeline energy over all 𝑄 GPUs in the pipeline is:

(1 + 𝛾) (𝑛𝑄 𝐶 ℓ 𝑢 𝑇 + 2 𝑢 𝑇 𝐻 𝑉 ) . 𝐹˜𝑔𝑄

(8)

𝐸 iter = 𝑇 iter ·

𝑄 ∑︁

𝑃𝑔idle +𝑀 · 𝑞

𝑞=1

The last stage computes its 𝑛𝑄 layers like any other stage, but additionally projects the hidden states to vocabulary logits: the extra term 2 𝑢 𝑇 𝐻 𝑉 is the cost of projecting each of the 𝑢 𝑇 tokens from hidden dimension 𝐻 to vocabulary size 𝑉 . The factor (1 + 𝛾) again accounts for the backward pass. Forward communication time. Stage 𝑞 sends activations forward to stage 𝑞 + 1 in 𝑆𝑞act (𝑢) , 𝐵˜ 𝑔𝑞 ,𝑔𝑞+1

comm 𝜏𝑞→𝑞+1 (𝑢) = 𝐿𝑔𝑞 ,𝑔𝑞+1 +

𝑆𝑞act (𝑢) . 𝐵˜ 𝑔𝑞 ,𝑔𝑞−1

4.5

(9)

(10)

4.6

(11)

Optimization

min

𝜆·

𝑄,{𝑛𝑞 },{𝑔𝑞 },𝑀

𝑄 ∑︁

𝑛𝑞 = 𝐿,

𝑛𝑞 ≥ 1 ∀𝑞,

(17)

𝑔𝑞 ∈ 𝐺, 𝑔𝑞 ≠ 𝑔𝑞 ′ ∀𝑞 ≠ 𝑞 ′, Mem𝑞train ≤ 𝑀˜ 𝑔𝑞 ∀𝑞, 𝑀∈Z ,

(12)

1 ≤ 𝑀 ≤ 𝐵.

(18) (19) (20)

𝑇 ref and 𝐸 ref are normalization constants, typically the time and energy of a baseline plan, chosen so the two terms are commensurable. At 𝜆 = 1 the problem reduces to pure time minimization (the regime of prior work). At 𝜆 = 0, it reduces to pure energy minimization. Intermediate values trace the Pareto frontier. The constraints enforce a valid layer partition (Eq. 17), distinct per-stage GPU assignment (Eq. 18), per-stage memory feasibility (Eq. 19), and a valid micro-batch count (Eq. 20).

actively computing. The remainder of the iteration, 𝑇 iter − 𝜏𝑔active , is spent idle in pipeline bubbles or waiting for up𝑞 stream/downstream stages. Per-GPU energy. Over a full iteration, GPU 𝑔𝑞 consumes 𝐸𝑔iter ≈ 𝑃𝑔idle · 𝑇 iter + (𝑃𝑔act − 𝑃𝑔idle ) · 𝜏𝑔active . 𝑞 𝑞 𝑞 𝑞 𝑞

(16)

𝑞=1

+

(𝑢)

𝑇 iter 𝐸 iter + (1 − 𝜆) · 𝑇 ref 𝐸 ref

subject to

We adopt a two-level power model for each GPU 𝑔𝑞 : idle power 𝑃𝑔idle and active power 𝑃𝑔act , where active power is 𝑞 𝑞 drawn while computing at effective capacity 𝐹˜𝑔𝑞 . Active time per iteration. Stage 𝑞 processes 𝑀 microbatches per iteration, so GPU 𝑔𝑞 spends comp

(15)

GeoTrain combines iteration time and energy into a single scalar objective via a user-tunable weight 𝜆 ∈ [0, 1]:

Energy Consumption Model

𝜏𝑔active ≈ 𝑀 · 𝜏𝑞 𝑞

Solution Representation

where 𝑄 is the number of pipeline stages, 1 ≤ 𝑄 ≤ 𝑁 ; {𝑛 1, . . . , 𝑛𝑄 } are the per-stage layer counts, with 𝑛𝑞 ≥ 1 and Í𝑄 𝑞=1 𝑛𝑞 = 𝐿, so every stage holds at least one layer and all 𝐿 layers are covered; and {𝑔1, . . . , 𝑔𝑄 } is the GPU assignment to stages, constrained to be disjoint, 𝑔𝑞 ≠ 𝑔𝑞 ′ ∀ 𝑞 ≠ 𝑞 ′ , so no GPU serves more than one stage. This representation jointly captures GPU subset selection (which 𝑄 of 𝑁 GPUs participate), layer partitioning, GPUto-stage assignment. The space grows rapidly with 𝑁 and 𝐿, motivating the algorithm in Sect. 5.

with the first and last stages omitting the incoming and outgoing communication terms, respectively. Substituting Eq. 11 into the 1F1B iteration-time (Eq. 1) yields 𝑇 iter (𝑀, 𝑢) as a function of the pipeline configuration. 4.4

(𝑢). (14)

𝑞=1

A plan specifies four tightly coupled decisions:  𝑄, {𝑛 1, . . . , 𝑛𝑄 }, {𝑔1, . . . , 𝑔𝑄 }},

Note that the reverse link (𝑔𝑞 , 𝑔𝑞−1 ) may have different bandwidth and latency than the forward link, which often ignored in prior work but critical in WANs with asymmetric touring. Per-stage service time. The total time stage 𝑞 spends on one micro-batch 𝑢 is comp comm comm 𝜏𝑞 (𝑢) = 𝜏𝑞 (𝑢) + 𝜏𝑞→𝑞+1 (𝑢) + 𝜏𝑞→𝑞−1 (𝑢),

comp

(𝑃𝑔act − 𝑃𝑔idle ) · 𝜏𝑞 𝑞 𝑞

This model exposes a fundamental tension in energyaware planning. Using fewer, more powerful GPUs reduces idle energy but increases per-GPU power draw, while spreading work across lower-power GPUs reduces peak draw but increases iteration time and idle energy. The optimum depends on workload, hardware, and the desired trade-off.

where 𝑆𝑞act (𝑢) = 𝜅 ·𝑢 ·𝑇 ·𝐻 is the activation tensor size. Transfer time is computed by adding latency to the microbatch’s activation size divided by the link effective bandwidth. Backward communication time. Gradients of the same size flow back from stage 𝑞 to stage 𝑞 − 1: comm 𝜏𝑞→𝑞−1 (𝑢) = 𝐿𝑔𝑞 ,𝑔𝑞−1 +

𝑄 ∑︁

(13)

The first term is the baseline draw over the full iteration; the second is the incremental cost of active computation. 5

Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi

𝐶 =[

ℓ1

ℓ2

ℓ3

ℓ4

ℓ5

ℓ6

ℓ7

ℓ8

1

1

1

2

2

3

3

3

𝑏1

Algorithm 1: GeoTrain’s Memetic Algorithm ]

Input: Population size 𝑃, tournament size 𝑘, max generations 𝐺 max , runtime limit 𝑇max , patience 𝑤, local-search rate 𝑟 ls , elite size 𝐸 Output: Best plan 𝐶 ∗ 1 population ← InitialPopulation(𝑃) ; 2 Evaluate fitness for each 𝐶 ∈ population ; ∗ 3 𝐶 ← arg max fitness; 𝑔stall ← 0; 𝑡 0 ← current time; 4 for generation ← 1 to 𝐺 max do 5 if current time −𝑡 0 ≥ 𝑇max then break// time limit; 6 if 𝑔stall ≥ 𝑤 then break // no improvement; 7 parents ← TournamentSelect(population, 𝑘); 8 offspring ← []; 9 for iter ← 1 to |parents| do 10 (𝑝 1, 𝑝 2 ) ← random pair from parents; 11 if 𝑝 1 .genes = 𝑝 2 .genes then skip; 12 offspring.Append(BoundaryCrossover(𝑝 1, 𝑝 2 )) // Alg. 2

𝑏2

𝑛 1 =3, GPU1 𝑛 2 =2, GPU2 𝑛 3 =3, GPU3

Figure 3. Chromosome encoding example

5

GeoTrain Design

GeoTrain solves the optimization problem using a Memetic Algorithm (MA) [13]. The design comprises a chromosome encoding, domain-specific genetic operators, and a capacityrebalancing local search. 5.1

Chromosome Encoding

We encode each candidate plan as an integer chromosome of length 𝐿: 𝐶 = [𝑥 1, 𝑥 2, . . . , 𝑥 𝐿 ] (21) where 𝑥 ℓ ∈ {1, . . . , 𝑁 } is the GPU assigned to layer ℓ. The three decision variables in Eq. 15 are implicitly encoded. The number of stages 𝑄 corresponds to groups of consecutive identical GPU IDs, the layer partition 𝑛𝑞 to group lengths, and the GPU assignment 𝑔𝑞 to the GPU ID of each group. Decoding requires a single linear scan. We fix the micro-batch size 𝑢 and thus the micro-batch count 𝑀 = 𝐵/𝑢, excluding 𝑢 from optimization; jointly optimizing 𝑢 is left to future work. Figure 3 shows encoding for 𝐿 = 8 and 𝑄 = 3. A chromosome is valid when (i) layers assigned to the same GPU form a contiguous group and (ii) per-stage memory bounds are satisfied (Eq. 19). 5.2

𝐶 best ← population.Get_Best(); foreach 𝑐 ∈ offspring do // Alg. 3 if 𝑐.fitness ≤ 𝐶 best .fitness then 𝑐 ← Mutate(𝑐) ;

13 14 15

foreach 𝑐 ∈ offspring do // Alg. 4 if rand() < 𝑟 ls then 𝑐 ← LocalSearch(𝑐) ;

16 17

offspring ← Deduplicate(offspring, population); population ← SteadyStateReplace(population, offspring, 𝐸); if best fitness improved then Update 𝐶 ∗ ; 𝑔stall ← 0; else 𝑔stall ← 𝑔stall + 1;

18 19

20 21 22 23

GeoTrain Algorithm (Main-Loop) Overview

24

Algorithm 1 presents the main MA loop. GeoTrain initializes a population and iteratively evolves it using evolutionary operators. The process stops when the runtime limit 𝑇max , no-improvement limit 𝑤, or generation limit 𝐺 max is reached. Initial population. To seed diversity, GeoTrain samples 𝑃 chromosomes. For each individual, it uniformly draws a stage count 𝑄 from 1, . . . , min(𝐿, 𝑁 ), selects 𝑄 − 1 split points between layers 1 and 𝐿 to define stage sizes, samples 𝑄 GPUs from 𝐺, and assigns them to stages in order. Infeasible chromosomes (Eq. 19) are rejected and resampled. This yields an initial population spanning stage counts and GPU subsets. Selection. GeoTrain uses tournament selection [12], which provides explicit control over selection pressure via the tournament size, whereas fitness-proportional methods such as roulette-wheel selection are sensitive to the scaling and distribution of fitness values [1]. Survivor replacement. Steady-state elitist replacement [28] is utilized. Each generation preserves the top-𝐸 individuals and replaces up to 𝑃 − 𝐸 worst incumbents with offspring. Duplicate offspring are discarded to maintain diversity. Fitness function. The fitness of a chromosome measures the quality of its encoded pipeline plan. GeoTrain decodes the

return 𝐶 ∗ ;

chromosome to recover the training plan, then computes iteration time 𝑇 iter (Eq. 1) and per-iteration energy 𝐸 iter (Eq. 14). These are combined into a single objective via Eq. 16; These are combined into a single objective via Eq. 16; fitness is its negative, so higher fitness indicates a better plan. 5.3

Stage-Boundary-Aware Crossover

Standard single-point crossover that randomly cuts a chromosome may disrupt the intrinsic stage structure of parent plans. We therefore propose a stage-boundary-aware best-cut crossover that restricts candidate cuts to stage boundaries and selects the fitness-maximizing cut (Alg. 2). Figure 4 illustrates the procedure on a six-layer example. Candidate cuts come from the union of stage boundaries across both parents, B1 ∪ B2 . Restricting cuts to stage boundaries preserves contiguous, well-balanced stages. Choosing the cut by fitness further steers recombination toward the most beneficial combination. Additionally, with probability 𝑝 fall , we evaluate all possible cuts to explore potentially better offspring beyond stage boundaries (Alg. 2, line 3). 6

GeoTrain : Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources

mut

(a)

1

1

1

2

2

2

(b)

1

1

2

2

3

3

(c)

1

1

2

2

3

3

1

1

2

2

2

2

boundary-shift

2

2

1

1

3

3

stage-swap

1

4

4

3

3

GPU-reassignment

mut

mut

active: {1,2,3}, unused: {4}

Figure 4. Stage-boundary-aware best-cut crossover (𝐿 = 6). Candidate cut positions B1 ∪ B2 = {2, 3} (dashed lines) are drawn from stage boundaries of 𝑃1 and 𝑃 2 . For each direction, all candidate cuts are evaluated; the cut 𝑏 ∗ yielding the highest offspring fitness is selected (red tick). 𝑐 1∗ takes 𝑃1 ’s prefix and 𝑃2 ’s suffix; 𝑐 2∗ takes 𝑃2 ’s prefix and 𝑃1 ’s suffix.

Algorithm 3: Contiguity-Preserving Mutation Input: Copied Offspring 𝑐 ′ : copy of 𝑐; Output: Mutated offspring 𝑐 ′ 1 Sample mutation type uniformly from {boundary-shift, stage-swap, GPU-reassignment}; 2 if boundary-shift then 3 B ← {𝑖 : 𝑐 ′ .genes[𝑖] ≠ 𝑐 ′ .genes[𝑖 − 1], 𝑖 ∈ 1..𝐿 − 1}; 4 𝑏 ← uniform sample from B; 5 𝑐 ′ .genes[𝑏 − 1] ← 𝑐 ′ .genes[𝑏]; 6 else if stage-swap then 7 (𝑠 1, 𝑠 2 ) ← TwoStagesSampled(𝑐 ′ ); 8 𝑐 ′ .genes[𝑠 1 .start : 𝑠 2 .end] ← 𝑐 ′ .genes[𝑠 2 ] ⊕𝑐 ′ .genes[𝑠 1 .end : 𝑠 2 .start] ⊕𝑐 ′ .genes[𝑠 1 ];

Input: Parents 𝑃1, 𝑃2 : 𝑃1 .genes ≠ 𝑃2 .genes; Fallback probability 𝑝 fall Output: Offspring list 1 B1 , B2 ← stage boundaries of 𝑃 1 , 𝑃 2 ; 2 if rand() < 𝑝 fall then 3 candidates ← {1, . . . , 𝐿 − 1} ; // full-scan fallback 4 else 5 candidates ← B1 ∪ B2 ; // boundary-restricted

else 𝑞 ← random stage in 𝑐 ′ ; 11 unused ← 𝐺 \ {all GPUs in 𝑐 ′ .genes}; 12 if unused = ∅ then return 𝑐; 13 𝑔′ ← uniform sample from unused; 14 Reassign all layers in stage 𝑞 to 𝑔′ ; 9

𝑏 1∗ ← None; 𝑓1∗ ← −∞; ∗ ∗ 7 𝑏 2 ← None; 𝑓2 ← −∞; 8 foreach 𝑏 ∈ candidates do 9 𝑐 1 ← 𝑃1 [:𝑏] ⊕ 𝑃2 [𝑏:]; // ⊕: concatenation 10 𝑐 2 ← 𝑃2 [:𝑏] ⊕ 𝑃1 [𝑏:]; 11 if 𝑐 1 .Validate() and Fitness(𝑐 1 ) > 𝑓1∗ then 12 𝑏 1∗ ← 𝑏; 𝑓1∗ ← fitness(𝑐 1 ); 6

14

10

if 𝑐 ′ .Validate() then return 𝑐 ′ ; 16 return 𝑐; 15

if 𝑐 2 .Validate() and Fitness(𝑐 2 ) > 𝑓2∗ then 𝑏 2∗ ← 𝑏; 𝑓2∗ ← fitness(𝑐 2 );

offspring ← []; if 𝑏 1∗ ≠ None then offspring.Add(𝑃1 [:𝑏 1∗ ] ⊕ 𝑃 2 [𝑏 1∗ :]); ∗ ∗ ∗ 17 if 𝑏 2 ≠ None then offspring.Add(𝑃 2 [:𝑏 2 ] ⊕ 𝑃 1 [𝑏 2 :]); 18 return offspring; 15

Stage-swap mutation. Select two pipeline stages uniformly at random and swap their positions (Alg. 3, lines 7-8). This reorders GPUs across pipeline positions, allowing heavy inter-stage boundaries to move onto faster links and alleviating communication bottlenecks on the asymmetric network, without changing the active GPU subset. GPU-reassignment mutation. Reassign a randomly selected stage 𝑞 to a uniformly sampled inactive GPU 𝑔′ ∈ 𝐺 \ 𝑔1, . . . , 𝑔𝑄 (Alg. 3, lines 10-14). This introduces a new GPU into the active subset while retiring 𝑔𝑞 , providing the MA’s primary mechanism for exploring alternative GPU subsets. If all 𝑁 GPUs are active, the operator is non-operational.

16

5.4

GPU4 replaces GPU2

Figure 5. Example of three contiguity-preserving mutation operators applied to a six-layer chromosome.

Algorithm 2: Stage-Boundary-Aware Best-Cut Crossover BoundaryCrossover(𝑝 1, 𝑝 2 )

13

1

Contiguity-Preserving Mutation

GeoTrain applies one of three mutation operators, selected uniformly at random, to each eligible offspring (Alg. 3; Fig. 5). All preserve layer contiguity by construction. Mutation targets offspring whose fitness does not exceed the population best, preserving elite solutions while diversifying others. If mutation violates the memory constraint (Eq. 19), the chromosome reverts to its parent. Boundary-shift mutation. Select a random stage boundary position 𝑏 (an index where 𝑥𝑏 ≠ 𝑥𝑏 −1 ) and set 𝑥𝑏 −1 ← 𝑥𝑏 , transferring the last layer of the left stage to the right stage (Alg. 3, lines 2-3). As the finest-grained move, it can produce a mutant that rebalances load between two adjacent stages, easing the slower one when it bottlenecks the pipeline.

5.5

Capacity-Rebalancing Local Search

The local search exploits the 1F1B cost model to refine offspring toward locally optimal load balance, making GeoTrain memetic rather than purely evolutionary. Figure 6 illustrates the effect on a concrete two-stage example. The search iteratively applies the convergence step in Alg. 4. At each step, the assigned GPUs are ranked by 𝑛𝑞 /𝐹˜𝑔𝑞 7

Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi

Before

1

1

1

1

2

6

After

Δ = 2 layers transferred

2

1

1

2

2

2

We evaluate GeoTrain via real-hardware experiments and larger-scale simulation to address these research questions: RQ1. Can GeoTrain find lower-cost pipeline configurations than state-of-the-art baseline on heterogeneous GPU environments? RQ2. How does GeoTrain scale as the number of GPUs and model layers grows? RQ3. How close are GeoTrain’s solutions to the globally optimal configuration, and for what environments is optimality recoverable? RQ4. How sensitive is GeoTrain under varying geo-distributed network conditions? RQ5. What trade-offs emerge between iteration time and energy consumption?

2

local search

Δ=2

𝜏2 = 1/𝐹˜1 𝜏1 = 4/𝐹˜1

𝜏1 = 2/𝐹˜1

𝜏2 = 2/𝐹˜1

←− load balanced −→

GPU1

Performance Evaluation

GPU2

←− bottleneck stage

GPU1 (𝐹˜1 , slow) GPU2 (2𝐹˜1 , fast)

Figure 6. Capacity-rebalancing local search example (𝐿 = 6, GPU2 is 2× faster than GPU1 ). Algorithm 4: Capacity-Rebalancing Local Search (one convergence step)

6.1

Experimental Setup

We implement GeoTrain in Python using the 1F1B iterationtime and two-level energy models (Sect. 4.4). We validate it on real hardware (Sect. 6.3), while other experiments use simulated heterogeneous GPU topologies to explore configurations beyond resource limits. Cluster environments. Real-hardware experiments run on a six-node cluster, each node equipped with a single GPU: an A100 (40GB), an A30 (24GB), three A40s (48GB each), and T4 (16GB). Nodes are provisioned with 4-60 vCPUs and 16240GB of host RAM. They are geographically located across different clouds (Nectar1 and Google Cloud Platform) and connected via a WireGuard VPN exhibiting heterogeneous communication characteristics with bandwidths ranging from 100 Mbps to 2 Gbps and latencies from sub-millisecond to 50 ms. By utilizing virtualization and shared tenancy on the cloud (e.g., NVIDIA’s Multi-Instance GPU technology), each GPU slices its full capacity to provide a certain effective FLOPS, which we have to profile before the planner runs. Regarding simulation, we use synthetic GPU topologies with various compute and memory capacity, and asymmetric network links, specific details are given per experiment. Model. We use a decoder-only Transformer in both real training and simulated studies. Specific configurations are summarized in Table A.1. In each case, the microbatch size 𝑢 and count 𝑀 satisfy 𝑢 · 𝑀 = 𝐵 and are fixed per experiment. Search parameters. GeoTrain uses population size 𝑃 =200, tournament size 𝑘 = 4, elite size 𝐸 = 2, local-search rate 𝑟 ls = 0.5, runtime limit 𝑇max = 120, s, no-improvement limit 𝑤 =100 generations, and maximum generations 𝐺 max =500.

Input: Copy of Chromosome 𝑐 ′ ← copy of 𝑐; Output: Refined chromosome 𝑐 ′ 1 repeat 2 Compute load ℓ𝑞 = 𝑛𝑞 /𝐹˜𝑔𝑞 for each stage 𝑞; 3 Sort stages descending by ℓ𝑞 ; 4 worst ← highest-load stage with transferable layers; 5 best ← lowest-load stage with available capacity; 6 if worst = None or best = None or 𝐹˜best ≤ 𝐹˜worst then break;
7 Δ ← min cap(best), 𝑛 worst ; 8 for 𝑖 ← 1 to Δ − 1 do 9 if (𝑛 best + 𝑖)/𝐹˜best ≥ (𝑛 worst − 𝑖)/𝐹˜worst then Δ ← 𝑖; 10 break; 𝑛 best ← 𝑛 best + Δ; 𝑛 worst ← 𝑛 worst − Δ; Rebuild 𝑐 ′ .genes from updated layer counts; ′ 13 until 𝑐 .genes unchanged; ′ 14 return 𝑐 ; 11

12

(which is proportional to each stage’s compute time) in descending order, exposing the current pipeline bottleneck ( Alg. 4, lines 2-5). The search identifies the worst stage (highest per-layer compute time with available memory headroom) and the best stage (lowest per-layer compute time with available memory headroom). If the best stage is strictly faster than the worst, meaning that transferring layers from worst to best would reduce the bottleneck, the search computes the optimal transfer count Δ (Alg. 4, lines 8-10): the maximum number of layers that can be transferred without making the original bottleneck stage faster than the least-loaded stage, thus maintaining a balanced pipeline. The chromosome is then rebuilt with the updated layer counts. The step repeats until unchanged, indicating convergence to a local optimum of the load-balancing objective. Interaction with the evolutionary operators. Following mutation, local search is applied to offspring with probability 𝑟 ls = 0.5 to balance exploitation and exploration. It refines selected offspring toward balanced load assignments while leaving others unchanged for evolutionary exploration.

6.2

Baselines

We compare against the following baselines: Metis. A deterministic search baseline adapted from state-ofthe-art Metis [30] to a PP-only , where each stage is restricted to a single GPU. The default (dlft.) variant always uses all GPUs, following the original method. The improved variant (impr.) exhaustively searches GPU subsets up to 7 GPUs and 1 https://ardc.edu.au/services/ardc-nectar-research-cloud/

8

GeoTrain : Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources

Table 2. Real-hardware scaling configurations. Avg. iteration times are measured over iterations 1-9, excluding warm-up (iteration 0). ± denotes the coefficient of variation.

Table 3. Iteration time vs. GPU count (𝐿 = 12, 𝐻 = 2048, fixed heterogeneity mix). Positive favours GeoTrain. impr. Metis

Pool

Cumulative training time (min)

2 4 6

L 16 24 32

GeoTrain (s)

dlft. Metis (s)

27.8 (±0.7%) 16.0 (±0.6%) 16.7 (±0.9%)

43.5 (±2.2%) 41.4 (±0.1%) 33.6 (±1.9%)

Gain

GPUs used

+36.3% +61.4% +50.1%

1 of 2 3 of 4 5 of 6

Cumulative Training Time on Real Hardware = time difference at each checkpoint

300 250

GeoTrain dflt. Metis

4 6 8 10 12

12 12 12 12 12

GeoTrain

Gain

(ms)

Best

Mean

CI

(%)

1979.8 1754.2 1354.5 1323.6 1099.5

1979.8 1750.2 1254.8 1233.4 996.6

1979.8 1750.2 1254.8 1233.4 996.6

±0.0 ±0.0 ±0.0 ±0.0 ±0.0

+0.0 +0.2 +7.4 +6.8 +9.4

1.4 h

150

impr. Metis

1.0 h 230m

100

172m

25 min

50 0

Layers

Table 4. Iteration time vs. model depth (GPUs = 12, 𝐻 = 2048, fixed heterogeneity mix). Positive favours GeoTrain.

1.9 h

200

GPUs

29m

54m

Iter 100

117m 55m

Iter 200

86m

Iter 300

114m

Iter 400

GeoTrain

Gain

GPUs

L

(ms)

Best

Mean

CI

(%)

12 12 12 12 12

8 12 24 32 48

807.4 1099.5 1822.1 2204.8 2191.8

801.9 996.6 1705.9 2199.3 1958.7

801.9 996.6 1705.9 2199.3 1998.6

±0.0 ±0.0 ±0.0 ±0.0 ±38.7

+0.7 +9.4 +6.4 +0.3 +10.6

Figure 7. Cumulative training time on real 6-GPU pool 6.4 uses random sampling beyond. The dlft. is applied in realhardware experiments and impr. for large-scale simulation, where the original Metis becomes infeasible. Both variants use the same analytical cost model as GeoTrain, with the energy disabled, since Metis does not optimize for energy. Brute-Force (BF) Search. Exhaustively enumerates pipeline configurations, including GPU subsets, orderings, and layer partitions, yielding the global optimum under the cost model and constraints. It is only practical for small setups. Branch-and-Bound (B&B). Incrementally constructs pipeline configurations and prunes partial solutions using iterationtime and feasibility bounds, preserving optimality while reducing enumeration cost. It is remains practical only for small setups.

Scalability

Moving beyond real hardware, we study RQ1, RQ2 through large-scale simulation, evaluating how GeoTrain’s advantage over impr. Metis varies with two problem dimensions: the number of available GPUs (Table 3) and model depth (Table 4). In each experiment, all other factors are held constant, so any change in the performance gap can be attributed solely to the parameter being varied. All reported times are iteration times based on the cost model. Gain % is (𝑇Metis −𝑇MA )/𝑇Metis ; positive values indicate that GeoTrain is faster. Across all ten configurations, GeoTrain matches or outperforms the improved state-of-the-art baseline. GPU count. The advantage over impr. Metis is non-monotonic with GPU pool size but becomes pronounced beyond 𝑁 = 8 (Table 3). For 𝑁 ≤ 6, both perform similarly (gap ≤ 0.2%), as limited heterogeneity leaves little optimization opportunity. For 𝑁 ≥ 8, GeoTrain outperforms impr. Metis by 6.8-9.4% through guided search that excludes slow devices and rebalances stages, while impr. Metis uses random subset sampling and proportional partitioning that can retain stragglers. Model depth. Model depth alone does not determine the performance gap (Table 4); the gain depends on how well the FLOPS-proportional Metis baseline balances the heterogeneous pipeline. When the proportional split is already near-balanced (𝐿 = 8, 32), the gap remains below 1%. Otherwise, proportional partitioning creates a straggler stage that GeoTrain improves through topology-aware exploration and rebalancing, yielding 6.4–10.6% gains. Convergence stability. Across all but the largest configuration (𝐿 = 48), GeoTrain converges to the same solution across five seeds, indicating stable convergence. Variance appears only at 𝐿 = 48 (±38.7ms), where the larger search space introduces multiple local optima. Nevertheless, GeoTrain

6.3 Real-Hardware Evaluation We evaluate RQ1 and RQ2 across three real-hardware configurations (Table 2), varying GPU availability and model depth while holding other factors fixed. Across all configurations, GeoTrain consistently outperforms dlft. Metis, reducing iteration time by 36.3%-61.4%. This gain stems from selectively excluding slower devices to mitigate stragglers. As shown in the GPUs used column, optimal plans use only a subset of GPUs, avoiding slowdowns from heterogeneous performance and communication. In contrast, dlft. Metis uses all available devices, making it more sensitive to stragglers. Additionally, the low coefficients of variation in Table 2 indicate stable execution, supporting the reliability of measured times and GeoTrain’s analytical model. As shown in Fig. 7, even modest per-iteration improvements can yield substantial savings over thousands of training iterations, reducing total training time by days or weeks. 9

Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi

Table 5. GeoTrain recovers the exact optimum on every computable scenario. Match counts seeds reaching optimum. N

L

𝑇 ∗ (ms)

GeoTrain (ms)

3 4 5 6

12 12 12 12

5291.165 1979.826 1878.830 1750.161

5291.165 1979.826 1878.830 1750.161

Optimality

Match

100% 100% 100% 100%

10/10 10/10 10/10 10/10

0.5-100,Gbps, respectively. We evaluate GeoTrain and impr. Metis across these settings, averaging 10 runs. Latency sensitivity. We fix cross-site bandwidth at 10 Gbps and sweep the base cross-site latency (Fig. 8b). Both planners exclude Site A at every level: its GPUs are weakest in compute and carry the highest cross-site penalty. From 0.1 ms to 10 ms, both select multi-site B+C plans with a single crosssite hop, and GeoTrain leads impr. Metis by 10.4% to 13.4% via better load balance. At the 50 ms peak (13.6%), impr. Metis abandons cross-site communication for a single-site plan on Site C (Fig. 9, top), as its proportional rule cannot justify a multi-site split. GeoTrain instead retains a balanced B+C plan (Fig. 9, bottom), where load balance offsets communication cost, cutting iteration time 1,996 to 1,757 ms. This highlights the value of topology-aware layer allocation. At 100 ms, both converge to Site C and the gap closes to 0%. Bandwidth sensitivity. We fix cross-site latency at 10 ms and sweep bandwidth from 0.5 to 100 Gbps (Fig. 8c). At low bandwidth, communication dominates and both planners consolidate onto a single site. As bandwidth increases, multisite execution becomes viable and GeoTrain exploits it faster than impr. Metis, opening its largest advantage (12.0% at 5 Gbps). Once bandwidth is abundant, communication is no longer a bottleneck, and the two converge. Overall, GeoTrain tracks network conditions more closely, shifting between single- and multi-site plans at finer granularity to capture gains missed by the baselines.

Table 6. Iteration time and planning time and vs. GPU count (𝐿 = 12). Exact solvers exhaust the 120 s budget at 𝑁 ≥ 8 (—). BF

B&B

GeoTrain

N

L

Space

iter. (ms)

plan.

iter. (ms)

plan.

iter. (ms)

plan.

2 4 6 8 10 12

12 12 12 12 12 12

24 5.4K 637K 43.8M 1.8B 40.9B

5839.7 1979.8 1750.2 — — —

0.0 s 0.4 s 16.6 s — — —

5839.7 1979.8 1750.2 — — —

0.0 s 0.3 s 15.9 s — — —

5839.7 1979.8 1750.2 1254.8 1233.4 996.6

13.3 s 15.9 s 18.8 s 22.8 s 32.5 s 29.4 s

consistently outperforms the baseline on average, demonstrating robustness under increased optimization complexity. 6.5

Optimality and Tractability

While Sect. 6.4 evaluates GeoTrain against impr. Metis, RQ1 and RQ3 require comparison with the global optimum. We therefore use two exact solvers as ground truth: Brute-Force (BF) and Branch-and-Bound (B&B). Correctness on heterogeneous scenarios. We first evaluate GeoTrain on four simulated heterogeneous scenarios (3 to 6 GPUs, four GPU types, 𝐿 = 12). As Table 5 shows, GeoTrain matches the exact optimum on all four scenarios across all ten seeds, yielding a 100.00% optimality (Opt.). Tractability and optimality We vary GPU count (𝑁 = 212, 𝐿 = 12) and compare BF, B&B, and GeoTrain (Table 6). The search space grows rapidly with 𝑁 , from just 24 configurations at 𝑁 = 2 to over 40 billion at 𝑁 = 12. where Ímin(𝑁 ,𝐿) 𝐿−1
𝑓 (𝑁 , 𝐿) = 𝑘=1 𝑘 −1 𝑃 (𝑁 , 𝑘), counts all valid layer partitions and GPU assignments. On all tractable instances (𝑁 ∈ {2, 4, 6}), GeoTrain matches the optimum, yielding a 100% optimality. Beyond 𝑁 = 6, BF and B&B fail to finish within the 120 s budget, while GeoTrain remains the only method that returns a plan. Together, these results demonstrate GeoTrain’s quality: it matches the verifiable optimum on tractable instances and remains effective beyond the reach of exact methods. 6.6

6.7

Energy-Time Trade-off

The objective in Eq. 16 enables GeoTrain to explore the Pareto frontier between iteration time and energy via 𝜆 ∈ [0, 1], from pure energy (𝜆 = 0) to pure time minimization (𝜆 = 1). To answer RQ5, we vary 𝜆 on a 14-GPU, 12-layer simulated setup, with five seeds per value. The GPU types are A30 (24 GB), A100 PCIe (40 GB), and A100 SXM (80 GB), with idle/active powers of 25/165 W, 40/250 W, and 60/400 W, respectively. We use rated TDP as active power and set idle power to approximately 15% of TDP, consistent with the ∼60 W idle power measured for the A100 SXM (80 GB) [3, 16– 18]. As the objectives differ in units, we normalize them using 𝑇ref = 870.4 ms and 𝐸 ref = 737.8 J from the impr. Metis plan. Table 7 reports the best (𝑇 iter, 𝐸 iter ) per 𝜆, while Fig. 10 shows the frontier and Fig. 11 shows how objectives and GPU subset size vary with 𝜆. All seeds converge to the same plan. Energy-saving (𝜆 ≤ 0.3). GeoTrain uses as few GPUs as possible, selecting A100 PCIe devices to minimise idle power from unused hardware. It notably avoids A30 GPUs despite their lower TDP, because A100 PCIe achieves better compute efficiency-lower peak power per effective FLOPS - so faster stage execution reduces total active-phase energy Balanced (𝜆 = 0.5). GeoTrain augments the A100 PCIe stages with two A100 SXM GPUs (4-GPU plan). Compared to the time-optimal plan (𝜆 =1), it is only 4.6% slower but uses

Network Sensitivity

We evaluate how planners adapt to geo-distributed networks and when topology-aware layer allocation yields the largest gains over the baseline (RQ4). Using a simulated three-site topology (Fig. 8a), we isolate cross-site effects: intra-site links provide 0.1,ms latency and 100,Gbps bandwidth, while cross-site links are asymmetric: A-B at 1.0×, B-C at 0.8× latency/1.2× bandwidth, and A-C at 1.5× latency/0.6× bandwidth. Base latency and bandwidth vary from 0.1-100,ms and 10

GeoTrain : Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources

Site C

2080Ti 3090

A30

A100

g0

g1

2080Ti

g2

g3

V100

g4 ×0.8 lat

×1.0 lat ×1.0 bw

0.1 ms 100 Gbps

g5

A30

0.1 ms 100 Gbps

×1.2 bw

Latency sweep (bw = 10 Gbps) 2000

g6

0.1 ms 100 Gbps

Bandwidth sweep (lat = 10 ms) 2000

impr. Metis GeoTrain (MA)

1900

Iteration time (ms)

Site B

Iteration time (ms)

Site A

1800 1700 1600

A C: ×1.5 lat / ×0.6 bw (slowest cross-site link)

0.1

1

5 10

1900 1800 1700 1600

impr. Metis GeoTrain (MA)

1500

50 100

0.5

Base cross-site latency (ms)

(a) Three-site topology

2

5

10

50 100

Base cross-site bandwidth (Gbps)

(b) Latency sweep (bw = 10 Gbps)

(c) Bandwidth sweep (lat = 10 ms)

Figure 8. Network sensitivity

GeoTrain

g5 (C)

3L

1,996 ms

9L

g3 (B)

g6 (C)

4L

g5 (C)

1L 1L

2.25 2.00 1.75 1.50 1.25 1.00

1,757 ms

6L

g4 (B)

0 2 4 6 8 10 12 Pipeline layers (bar width proportional to layers per stage)

0.0 0.1

Figure 9. Plans found at 50 ms cross-site latency.

𝝀

𝑻 iter (ms) 1,976.2 1,232.7 1,041.7 839.3 802.3 802.3 802.3

𝑬 iter (J)

𝚫𝑻

𝚫𝑬

519.4 +146.3% −25.0% 541.4 +53.7% −21.8% 571.1 +29.8% −17.6% 651.5 +4.6% −6.0% 692.8 — — 692.8 — — 692.8 (baseline) (baseline)

impr. Metis (reference): — 870.4 737.8

+8.5%

Energy per iteration E iter (J)

impr. Metis (ref) GeoTrain ( sweep) =1.0 (time-only) Knee ( =0.1) Pareto frontier

impr. Metis (ref, dominated)

700

=0.7/0.9/1.0 (time-optimal) 650

=0.5

600

=0.3

500

Knee ( =0.1) 21.8% energy +53.7% time vs. =1.0 =0.1

550 GeoTrain ( 0.5) dominates Metis (lower T AND lower E)

800

1000

1200

=0.0 1400

0.8 0.3

0.5

0.7

0.9 1.0

(time weight)

0.7

4 2

2 1

0 .0 .1 0 0

0.3

0.5

0.7

0.9 1.0

(time weight)

Time-optimal (𝜆 ≥ 0.7). All five seeds converge to an identical 5-GPU plan for 𝜆 ∈ [0.7, 1.0], indicating that the timeoptimal configuration is a stable fixed point and that increasing 𝜆 beyond 0.7 does not change the solution. GeoTrain strictly dominates Metis. The impr. Metis plan (870.4,ms, 737.8,J) is dominated by all GeoTrain plans with 𝜆 ≥ 0.5. At 𝜆 = 0.5, GeoTrain reduces iteration time and energy by 3.6% and 11.7%, respectively. At 𝜆 = 1.0, it achieves the fastest configuration (802.3,ms) while still reducing energy by 6.1%. The knee gives the best return on energy. The maximumperpendicular-distance criterion [23] places the knee at 𝜆 = 0.1, whose two-GPU plan cuts iteration time by 37.6% over the single-GPU energy-only plan (𝜆 =0) for only 4.2% more energy. It is the recommended point for energy-constrained deployments needing high throughput. Those requiring nearoptimal speed can use 𝜆 =0.5, within 4.6% of the time-only optimum while saving 6.0% energy.

+6.5% 3×A100 SXM, 2×A100 PCIe

750

0.9

sweep =1.0 (time-only) Knee ( =0.1) 5 5 5 4 3

6

GPUs selected 1×A100 PCIe 2×A100 PCIe 3×A100 PCIe 2×A100 SXM, 2×A100 PCIe 2×A100 SXM, 3×A100 PCIe 2×A100 SXM, 3×A100 PCIe 2×A100 SXM, 3×A100 PCIe

★ Knee of the frontier (maximum perpendicular distance from the line joining the extreme points).

800

(b) Plan GPU count vs.\

Figure 11. (a) Normalised iteration time (𝑇 iter /𝑇ref , blue) and energy (𝐸 iter /𝐸 ref , red) vs. 𝜆; dotted lines mark the impr. Metis reference (= 1). (b) GPUs in the winning plan vs. 𝜆.

Table 7. Energy-time Pareto frontier (14 GPUs, 12 layers, 𝐻 =2048). Δ𝑇 and Δ𝐸 are deviations from the 𝜆 =1.0 plan.

0.0 0.1★ 0.3 0.5 0.7 0.9 1.0

(a) Objectives vs.\ T iter / Tref 1.0 knee E iter / Eref

Norm.\ energy GPUs in winning plan

impr. Metis

g6 (C)

Site B (A30)

Norm.\ iteration time

Site C (A100 / V100)

1600

1800

Iteration time T iter (ms)

2000

7

Figure 10. Energy-time Pareto frontier of GeoTrain (14 GPUs, 12 layers). Points show the best plan across five seeds for each 𝜆. Dashed line connects the frontier.

Conclusion and Future Work

We presented GeoTrain, an automatic parallelism system for foundation model training on heterogeneous, geo-distributed GPUs. By jointly optimizing GPU selection, model partitioning, and placement under compute, memory, power, and network constraints, GeoTrain efficiently explores the combinatorial search space using an optimized memetic algorithm.

6.0% less energy, making it a strong energy-aware operating point with near-optimal speed. 11

Bui-Khanh Hoang, TianZhang He, Mohammad Goudarzi

Across real and simulated environments, GeoTrain outperforms state-of-the-art baselines, reducing iteration time and energy by up to 61.4% and 11.7%, respectively, while matching the verifiable optimum on small instances and scaling beyond exact methods. These results show that topologyand energy-aware planning can provide a practical and scalable approach to distributed foundation model training. In future work, we will extend GeoTrain to support additional parallelism techniques within pipeline stages, adapt to dynamic network and resource conditions, incorporate additional practical objectives, and investigate larger-scale production deployments.

[12] Brad L. Miller and David E. Goldberg. 1995. Genetic Algorithms, Tournament Selection, and the Effects of Noise. Complex Systems 9, 3 (1995), 193–212. [13] Pablo Moscato and Carlos Cotta. 2003. A Gentle Introduction to Memetic Algorithms. In Handbook of Metaheuristics, Fred Glover and Gary A. Kochenberger (Eds.). International Series in Operations Research & Management Science, Vol. 57. Springer, Boston, MA, 105– 144. doi:10.1007/0-306-48056-5_5 [14] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, et al. 2019. PipeDream: Generalized Pipeline Parallelism for DNN Training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles (SOSP 19). 1–15. [15] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. 2021. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC ’21). 1–14. doi:10.1145/3458817. 3476209 [16] NVIDIA Corporation. 2020. NVIDIA A100 40GB PCIe GPU Accelerator Product Brief. https://www.nvidia.com/content/dam/en-zz/Solutions/ Data-Center/a100/pdf/A100-PCIE-Prduct-Brief.pdf. States the GPU operates unconstrained up to its TDP level of 250 W. Accessed 202606-11. [17] NVIDIA Corporation. 2021. NVIDIA A100 Tensor Core GPU Datasheet. https://www.nvidia.com/content/dam/en-zz/Solutions/DataCenter/a100/pdf/nvidia-a100-datasheet-nvidia-us-2188504web.pdf. Max Thermal Design Power: 250/300 W (PCIe), 400 W (SXM). Accessed 2026-06-11. NVIDIA A30 Tensor Core GPU [18] NVIDIA Corporation. 2021. Datasheet. https://www.nvidia.com/content/dam/en-zz/Solutions/ data-center/products/a30-gpu/pdf/a30-datasheet.pdf. Max Thermal Design Power: 165 W. Accessed 2026-06-11. [19] Jay H. Park, Gyeongchan Yun, Chang M. Yi, et al. 2020. HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism. In 2020 USENIX Annual Technical Conference (USENIX ATC 20). 307–321. [20] David Patterson, Joseph Gonzalez, Quoc Le, et al. 2021. Carbon Emissions and Large Neural Network Training. arXiv preprint arXiv:2104.10350. [21] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC 20). [22] Max Ryabinin, Tim Dettmers, Michael Diskin, and Alexander Borzunov. 2023. SWARM Parallelism: Training Large Models Can Be Surprisingly Communication-Efficient. In Proceedings of the 40th International Conference on Machine Learning (ICML ’23) (ICML’23). JMLR.org, Honolulu, Hawaii, USA. [23] Ville Satopää, Jeannie Albrecht, David Irwin, and Barath Raghavan. 2011. Finding a “Kneedle” in a Haystack: Detecting Knee Points in System Behavior. In 2011 31st International Conference on Distributed Computing Systems Workshops (ICDCSW). IEEE, 166–171. doi:10.1109/ ICDCSW.2011.20 [24] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, et al. 2019. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv:1909.08053. [25] Foteini Strati, Zhendong Zhang, George Manos, Ixeia Sánchez Périz, Qinghao Hu, Tiancheng Chen, Berk Buzcu, Song Han, Pamela Delgado, and Ana Klimovic. 2025. Sailor: Automating Distributed Training over Dynamic, Heterogeneous, and Geo-distributed Clusters. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles

References [1] Tobias Blickle and Lothar Thiele. 1996. A Comparison of Selection Schemes Used in Evolutionary Algorithms. Evolutionary Computation 4, 4 (1996), 361–394. doi:10.1162/evco.1996.4.4.361 [2] Tom Brown, Benjamin Mann, Nick Ryder, et al. 2020. Language Models are Few-Shot Learners. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 33. 1877–1901. [3] Qinghao Hu, Zhisheng Ye, Zerui Wang, Guoteng Wang, Meng Zhang, Qiaoling Chen, Peng Sun, Dahua Lin, Xiaolin Wang, Yingwei Luo, Yonggang Wen, and Tianwei Zhang. 2024. Characterization of Large Language Model Development in the Datacenter. In 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24). USENIX Association, Santa Clara, CA, USA, 709–729. https: //www.usenix.org/conference/nsdi24/presentation/hu [4] Yanping Huang, Youlong Cheng, Ankur Bapna, et al. 2019. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 32. [5] Insu Jang, Zhenning Yang, Zhen Zhang, Xin Jin, and Mosharaf Chowdhury. 2023. Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templates. In Proceedings of the 29th Symposium on Operating Systems Principles (SOSP ’23). Association for Computing Machinery, 382–395. doi:10.1145/3600006.3613152 [6] Byungsoo Jeon, Mengdi Wu, Shiyi Cao, Sunghyun Kim, Sunghyun Park, Neeraj Aggarwal, Colin Unger, Daiyaan Arfeen, Peiyuan Liao, Xupeng Miao, Mohammad Alizadeh, Gregory R. Ganger, Tianqi Chen, and Zhihao Jia. 2025. GraphPipe: Improving Performance and Scalability of DNN Training with Graph Pipeline Parallelism (ASPLOS ’25). Association for Computing Machinery, New York, NY, USA, 557–571. doi:10.1145/3669940.3707220 [7] Jared Kaplan, Sam McCandlish, Tom Henighan, et al. 2020. Scaling Laws for Neural Language Models. arXiv preprint arXiv:2001.08361 (2020). [8] Diederik P. Kingma and Jimmy Ba. 2015. Adam: A Method for Stochastic Optimization. In Proceedings of the 3rd International Conference on Learning Representations (ICLR). [9] Dacheng Li, Hongyi Wang, Eric Xing, and Hao Zhang. 2022. AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awareness. Advances in Neural Information Processing Systems (NeurIPS) 35 (2022), 6630–6639. [10] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. 2020. PyTorch Distributed: Experiences on Accelerating Data Parallel Training. Proceedings of the VLDB Endowment 13, 12 (2020), 3005–3018. doi:10.14778/3415478.3415530 [11] Xupeng Miao, Yujie Wang, Youhe Jiang, Chunan Shi, Xiaonan Nie, Hailin Zhang, and Bin Cui. 2022. Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism. Proc. VLDB Endow. 16, 3 (Nov. 2022), 470–479. doi:10.14778/3570690.3570697 12

GeoTrain : Energy-Aware Automatic Parallelism for Foundation Model Training on Heterogeneous Geo-Distributed Resources

(SOSP ’25). Association for Computing Machinery, Seoul, Republic of Korea, 204–220. doi:10.1145/3731569.3764839 [26] Emma Strubell, Ananya Ganesh, and Andrew McCallum. 2019. Energy and Policy Considerations for Deep Learning in NLP. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL). 3645–3650. [27] Zhenbo Sun, Huanqi Cao, Yuanwei Wang, Guanyu Feng, Shengqi Chen, Haojie Wang, and Wenguang Chen. 2024. AdaPipe: Optimizing Pipeline Parallelism with Adaptive Recomputation and Partitioning. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS ’24), Volume 3. Association for Computing Machinery. doi:10.1145/ 3620666.3651359 [28] Gilbert Syswerda. 1991. A Study of Reproduction in Generational and Steady-State Genetic Algorithms. In Foundations of Genetic Algorithms, Gregory J. E. Rawlins (Ed.). Vol. 1. Morgan Kaufmann, 94–101. doi:10. 1016/B978-0-08-050684-5.50009-4 [29] Hugo Touvron, Thibaut Lavril, Gautier Izacard, et al. 2023. LLaMA: Open and Efficient Foundation Language Models. arXiv preprint arXiv:2302.13971. [30] Taegeon Um, Byungsoo Oh, Minyoung Kang, et al. 2024. Metis: Fast Automatic Distributed Training on Heterogeneous GPUs. In 2024 USENIX Annual Technical Conference (USENIX ATC 24). 563–578. [31] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention Is All You Need. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 30. 5998–6008. [32] Jue Wang, Yucheng Lu, Binhang Yuan, Beidi Chen, Percy Liang, Christopher De Sa, Christopher Re, and Ce Zhang. 2023. CocktailSGD: Fine-tuning Foundation Models over 500Mbps Networks. In Proceedings of the 40th International Conference on Machine Learning (ICML ’23) (ICML’23). JMLR.org, Honolulu, Hawaii, USA. [33] Jie You, Jae-Won Chung, and Mosharaf Chowdhury. 2023. Zeus: Understanding and Optimizing GPU Energy Consumption of DNN Training. In 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23). 119–139. [34] Binhang Yuan, Yongjun He, Jared Quincy Davis, Tianyi Zhang, Tri Dao, Beidi Chen, Percy Liang, Christopher Re, and Ce Zhang. 2022. Decentralized Training of Foundation Models in Heterogeneous Environments. In Proceedings of the 36th International Conference on Neural Information Processing Systems (NeurIPS ’22). Curran Associates Inc., New Orleans, LA, USA. [35] Lianmin Zheng, Zhuohan Li, Hao Zhang, et al. 2022. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). 559–578. [36] Zan Zong, Minkun Guo, Mingshu Zhai, Yinan Tang, Jianjiang Li, and Jidong Zhai. 2025. Training Large Models on Heterogeneous and GeoDistributed Resource with Constricted Networks. Big Data Mining and Analytics 8, 4 (2025), 966–980. doi:10.26599/BDMA.2025.9020031

A

Hidden size 𝐻 Depth 𝐿 (layers) Microbatch size 𝑢 Microbatches 𝑀 Batch size 𝐵 Sequence length 𝑇 Vocabulary size Precision

Experimental Configurations

Table 1. Transformer configuration details in real-world and simulated experiments 13

Real training

Simulated studies

1024 16-32 16 64 1024 128 50,258 FP32

2048 8-48 8 8 64 2048 50,258 FP16

Related documents

• GeoFitBench-LLM: A Benchmark for Assessing Large Language Models' Capabilities in Geospatial Data Requirements Elicitation for Geo-Analytic Tasks#1296676
• Accelerating Sharded Data Parallelism at Scale with Federated Learning#1034123
• Accelerating Sharded Data Parallelism at Scale with Federated Learning#978390
• No One Knows the State of the Art in Geospatial Foundation Models#999791
• A Composition-Aware Pretraining Framework for Geospatial Foundation Models#622589