ETH: Data-Balanced Heterogeneous Pipeline Parallelism for Training Large Models

1 Introduction

• Heterogeneous layer assignment can leave devices underutilized: memory constraints limit high-compute devices, while coarse assignments add 16% pipeline bubbles.Conventional balancing can leave up to 32% of computation power idle; with only five layers per GPT3-13B stage, one layer changes workload by 20%.Source: §1 · pp. 1–2
• HPP uses independent heterogeneous pipelines with uneven local batches for fine-grained workload balance, while assigning layers and parallelism strategies within each pipeline.Memory use depends on micro-batch size, so data balancing need not align device memory capacity with computational throughput.Source: §1 · p. 2
• ETH jointly adjusts local batch sizes across pipelines and stage workloads within pipelines, using block-wise gradient aggregation to reduce communication overhead.Its search optimizations include pipeline caching, throughput prediction, priority-based ordering, and dynamic programming.Source: §1 · p. 2
• ETH achieves up to 2× faster training and reduces configuration generation time by up to 56× compared with existing training systems.The system extends Megatron-LM and was evaluated across diverse accelerators and scales.Source: §1 · p. 2

2.1 Distributing Training Parallelism

• Pipeline parallelism assigns consecutive-layer stages and processes micro-batches; data parallelism distributes disjoint input subsets, while tensor parallelism splits layer tensors across accelerator groups.Local batch size affects execution time but not memory consumption; micro-batch size affects memory overhead.Source: §2.1 · pp. 2–3

2.2 Limitations of Existing Works

• At 350M→175B, HexiScale ratios rise 1.34→1.92 and Metis 1.27→1.87; unlimited-memory Metis stays 1.09–1.13, while layer balancing leaves up to 32% computation idle.Even without memory constraints, latency ratios exceed ideal by 9–13%, resulting in 16% additional pipeline bubbles.Source: §2.2 · p. 3

2.3 Opportunities

• As mismatch rises from 1.0 to 1.9, GPT3-13B Metis’ latency ratio rises from 1.11 to 1.87, while data balancing remains between 1.09 and 1.17.DeepSeekV3’s global batch size exceeds its layer count by ≥251×, supporting finer workload granularity than layer assignment.Source: §2.3 · p. 3

Data-balanced heterogeneous pipeline parallelism

• The proposed data-balanced heterogeneous training paradigm motivates HPP as a way to balance work across heterogeneous pipelines.The passage introduces the paradigm but does not provide further detail here.Source: §3 · p. 3

Challenge 1: The vast search space for an optimal data-balanced HPP

• HPP configuration search can cost more than its training-time savings, and generation time grows superlinearly with device count.For data-balanced HPP, the generation-to-computation ratio ranges from 0.8 to 4905.Source: §3 · p. 4

Challenge 2: High communication overhead of asymmetric All-Reduce

• Devices in the same data-parallel group can have inconsistent gradient addresses and sizes, preventing direct All-Reduce.Device-wise padding adds memory and communication overhead, while tensor-wise All-Reduce creates a separate connection per tensor and can cause multi-stream bottlenecks.Source: §3 · p. 4, Figure 5 · p. 4

4 System Design

• ETH’s system design combines heterogeneous pipeline configuration generation, hierarchical workload balancing, independent pipeline execution, and block-wise gradient synchronization.Configuration-search optimizations include caching, early skipping of suboptimal groups, throughput-prioritized ordering, and dynamic programming.Source: §4 · p. 5, §4.1 · p. 5

4.2.1 Pipeline generation

• Pipeline generation treats pipelines as unordered sets, so configurations with the same device-group composition are generated only once.For each pipeline count, ETH generates distinct groupings, from one pipeline containing all devices to one pipeline per device.Source: §4.2.1 · p. 5

4.2.2 Stage device groups

• Stage-device groups preserve stage order because activation overhead is higher in earlier stages under the 1F1B schedule.Thus, reversing the order of two device groups represents a distinct configuration.Source: §4.2.2 · p. 5

4.2.3 Hierarchical workload balance

• ETH uses a three-phase heuristic: optimize stages with global-batch initialization, allocate local batches across pipelines, then refine stage configurations for those batches.The heuristic seeks low predicted iteration time but does not guarantee a globally optimal joint configuration.Source: §4.2.3 · pp. 5–6
• Stage search selects TP and DP degrees, layer counts, recomputation, and micro-batch sizes to reduce pipeline latency.Its latency objective combines first-micro-batch stage times with the slowest-stage latency multiplied by the remaining micro-batches.Source: §4.2.3 · p. 5
• Stage-level bottleneck adjustments respond differently to latency and memory limits, changing operator placement, parallelism degrees, micro-batch size, or recomputation.For latency bottlenecks, ETH may increase DP and decrease TP; for memory bottlenecks, it may increase TP and decrease DP.Source: §4.2.3 · p. 6
• ETH allocates local batches iteratively to the pipeline whose increment would cause the smallest increase in maximum completion time, stopping when the aggregate batch constraint is met.Each pipeline’s stage configuration is then adjusted using its assigned local batch size.Source: §4.2.3 · p. 6

4.3 Execution of Heterogeneous Pipelines

• Block-wise All-Reduce partitions and combines gradients to align blocks across devices, then fuses blocks sharing a data-parallel group.This maintains address and size consistency while reducing the number of All-Reduce operations.Source: §4.3 · pp. 6–7, Figure 7 · p. 6

5.1 Caching Existing Pipeline Configurations

• Caching stores a pipeline’s best configuration during its first optimization and reuses it when that pipeline recurs in other device groups.Cached settings include TP/DP degrees, layer count, recomputation operations per stage, and micro-batch size.Source: §5.1 · p. 7

5.2 Search via Early Skipping

• Early skipping prunes a device group when the smaller of its deterministic and approximate throughput bounds falls below the best configuration found so far.The deterministic bound uses device-capacity dominance checked with bipartite matching; the approximate bound scales cached throughput by computation and memory ratios with a fault-tolerance factor.Source: §5.2 · p. 7

5.3 Optimization for the Search Order

• ETH orders candidate groups using normalized inter- and intra-pipeline computation variance, prioritizing lower-cost groups within each pipeline count.Across sampled configurations, throughput fell as inter-pipeline variance increased; lower intra-pipeline variance tracked higher throughput, but its impact diminished as model depth increased.Source: §5.3 · pp. 7–8

5.4 Search via Dynamic Programming

• As viable configurations become cached, pipeline-combination growth becomes the search bottleneck; dynamic programming addresses this exponential complexity.For each pipeline count, it selects a split maximizing throughput between a configuration of the first devices and one additional pipeline.Source: §5.4 · p. 8
• Dynamic programming reduces HPP configuration search from exponential to polynomial time, with O(N^3) complexity when the maximum pipeline count equals device count.Afterward, ETH data-balances pipelines and refines each pipeline’s stage configuration.Source: §5.4 · p. 8

6 System Implementation

• ETH adds approximately 9,000 lines of Python to Megatron-LM, including stage-specific recomputation support, adaptive balancing-time allocation, and ZeRO-1/ZeRO-2 sharding.The additions also rewrite the generator to expand the search space using Aceso’s stage-specific recomputation.Source: §6 · p. 8
• Before searching, ETH profiles each accelerator type and benchmarks operator latency, communication, parameter and activation sizes, and reserved memory.Profiling covers candidate micro-batch sizes and TP degrees.Source: §6 · pp. 8–9
• When a stage exceeds device memory capacity during planning, ETH prioritizes recomputing high-memory, low-recompute-cost operators until the stage fits.The resulting plan discards selected activations after forward passes and recomputes them for gradient calculations.Source: §6 · p. 9

7.1 Experimental Setup

• Experiments cover GPT-3 models from 350M to 175B parameters and LLaMA-2 70B across Ascend 910B and NVIDIA A100/RTX 3090 clusters.The main evaluations vary 2–8 accelerators per node and 2–128 nodes, using sequence length 2048 and global batch size 1024.Source: §7.1 · p. 9, Table 2 · p. 9

7.2 End-to-End Performance

• ETH achieves 1.36–2.00× speedups over HexiScale across B2/B3 and B2/B4 clusters, and 1.20–1.84× over Metis; it also reduces search time versus both baselines.For 175B models, configuration generation is 38.22–42.99× faster than Metis.Source: §7.2 · p. 9, Figure 9 · p. 10

7.3 Computation and Communication Time

• For GPT-3 13B, ETH computes 30% faster than Metis and 40% faster than HexiScale, while reducing pipeline bubbles to 75% and 60% of their respective levels.Block-wise All-Reduce is 1.38–1.71× faster than tensor-wise and 1.56–2.74× faster than device-wise All-Reduce.Source: §7.3 · p. 10, Figure 10 · p. 10

7.4 System Scalability

• On B2/B4 clusters, ETH reaches up to 1.86× and 1.83× speedups over HexiScale and Metis, and its advantage grows with global batch size and sequence length.ETH trains at every evaluated device scale, while HexiScale encounters OOM at 16, 32, and 48 devices; ETH supports sequence lengths up to 8192.Source: §7.4 · p. 10, Figure 11 · p. 11

7.5 Evaluation of Key Components

• Disabling stage-specific recomputation raises peak memory by 1.8%–5.7%; runtime and memory prediction errors remain within 7.4%.Memory predictions consistently overestimate requirements.Source: §7.5 · p. 11, Figure 12 · p. 11

7.6 The Effectiveness of Search Optimization

• In the evaluated GPT-3 13B B2/B4 workload, ETH’s search optimizations reduced search cost while preserving generated-configuration iteration time; near-constant search cost and training time occurred when α exceeded 1.5.At α=0.5, search cost was 127.64 s and generated-configuration iteration time was 38.14 s.Source: §7.6 · p. 11, Figure 13 · p. 12

7.7 Model Convergence Performance

• Across 2,000 LLaMA-2 70B training iterations on Enwiki, ETH and the baselines show nearly identical loss curves, with differences of approximately 0.1%.ETH also produces consistent loss curves across the evaluated heterogeneous and homogeneous clusters.Source: §7.7 · p. 12

Search quality for HPP

• Aceso’s bottleneck-driven local stage adjustments can reach local optima, so ETH does not guarantee globally optimal HPP configurations.ETH currently searches PP, TP, and DP; adding CP or EP would require new workload, communication, memory, placement, and synchronization models.Source: §8 · p. 12

Other pipeline schedules

• ETH’s latency and memory models assume 1F1B scheduling, so other pipeline schedules require new cost models and execution-order validation.Alternative schedules can change pipeline bubbles and activation lifetimes, affecting feasibility and workload balance.Source: §8 · p. 12

• Prior work explores device placement, tensor parallelism, recomputation, and hierarchical parallelism; heterogeneous systems address workload mapping, layer or batch balancing, and cluster allocation.ETH jointly optimizes local batch size and stage configurations.Source: §9 · p. 12

• ETH jointly adjusts local batch sizes across pipelines and rebalances stage workloads within each pipeline, while using block-wise All-Reduce to reduce communication overhead.The paper identifies expansion of the available configuration space for heterogeneous parallel training as ETH’s key contribution.Source: §10 · p. 12

A Visualization of Training Configuration

• Table 3 visualizes the strategies generated by Metis, HexiScale, and ETH; its caption defines #L as the number of layers in each stage.The caption also expands Qkv, Sco, AS, and AD as attention-related labels.Source: Table 3 · p. 19
• Metis uses one pipeline with layer balancing, while HexiScale cannot use data balancing and fixes the data-parallel degree at 1 within each stage.HexiScale’s resulting configuration is a single pipeline with a data-parallel degree of 2 per stage.Source: A Visualization of Training… · p. 19
• ETH uses two data-parallel pipelines in the example: B4 processes 344 samples at 4 samples per micro-batch, and B2 processes 680 at 8.ETH also permits stage-specific data parallelism or tensor parallelism, layer counts, and recomputation.Source: A Visualization of Training… · p. 19

B Model Convergence Performance

• Loss curves differ by approximately 0.1% across the reported settings.The comparison includes ETH and baselines on LLaMA-2 70B trained on Enwiki for 2,000 iterations, with B2/B3 heterogeneous clusters and a B3 homogeneous cluster.Source: B Model Convergence Performance · p. 19
• The reported curves are nearly identical for ETH and the baselines on LLaMA-2 70B trained on Enwiki for 2,000 iterations.The experiments include B2/B3 heterogeneous clusters and a B3 homogeneous cluster.Source: B Model Convergence Performance · p. 19
• The results are reported to confirm that ETH improves system efficiency without affecting model convergence.Source: B Model Convergence Performance · p. 19