Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism

• Scaling Transformer parameters increases computational and memory costs, while parallelism choices yield different memory use, communication overhead, and execution efficiency.Source: §1 · p. 1
• Existing automatic approaches support limited parallelism dimensions or rely on expert-designed strategies, leaving a larger multi-dimension search space largely unexplored.Source: §1 · pp. 1–2
• Galvatron combines DP, SDP, TP, and PP, using decision-tree decomposition and pruning followed by dynamic programming to search hybrid strategies.Its cost estimation accounts for computation and communication overlap slowdown.Source: §1 · p. 2
• Galvatron achieves up to 338% and 55% throughput speedups over state-of-the-art pure and hybrid parallelism methods, respectively.Source: §1 · p. 2

2.1 Transformer Models

• Transformer layers primarily contain self-attention and feed-forward modules, whose dense tensor operations incur substantial computation and memory costs.Source: §2.1 · p. 2
• NLP Transformers include encoder-only, decoder-only, and encoder-decoder forms, while ViT uses image patches and Swin uses hierarchical shifted-window attention.The passage notes distinct computation and memory characteristics across NLP architectures and uneven computation and memory across Swin layers.Source: §2.1 · p. 2

2.2 Parallelism in Distributed Training

• Data parallelism distributes samples while retaining a model replica per worker; ZeRO partitions model states while retaining the data-parallel computation process.The passage also identifies ZeRO as FSDP in FairScale and cites ZeRO [37].Source: §2.2 · p. 2
• Model parallelism partitions computation: TP divides Transformer modules across devices, while PP assigns layer sequences to pipeline stages and transfers boundary activations.Megatron-LM uses TP [29], and GPipe introduced PP [13].Source: §2.2 · p. 2
• Prior automatic systems combine data parallelism with tensor or pipeline parallelism, but generally explore at most one model-parallel dimension alongside data parallelism.The cited discussion says task-parallel approaches address simultaneous training tasks rather than parallel training of a single task.Source: §2.2 · pp. 2–3
• Industrial Transformer training examples motivate automatic systems that cover more parallelism decisions without depending on strong expert tuning experience.Source: §2.2 · p. 3

3.1.1 Two-GPU Example.

• For each layer, the search must choose among DP, SDP, PP, and TP strategies with different memory costs and communication overheads.Backpropagation requires forward activations to remain in device memory until gradients are calculated.Source: §3.1.1 · p. 3

3.2 Decision-tree-based Decomposition

• The decision tree represents arbitrary combinations of parallelism without overlap, using PP partitions and device-group sizes to organize the search.Its rules prohibit repeating a parallelism at multiple levels and use non-leaf degrees from {2, 4, 8, ...}.Source: §3.2 · p. 4
• After pruning DP+SDP combinations, Galvatron retains 22 candidate hybrid strategies for a single layer on 8 GPUs.The decision tree initially produces 34 candidates, reducing a space of hundreds by an order of magnitude.Source: §3.2 · p. 4

3.1 Search Space Analysis

• Galvatron seeks the highest-throughput layer-level plan for a given model and device memory capacity.The search space grows exponentially with model layers, making brute-force enumeration infeasible within a feasible time budget.Source: §3.1.2 · p. 3, §3.2 · p. 3, §3.3 · p. 4

3.3 Parallelism Optimization

• Dynamic programming minimizes stage execution time under memory budget E, accounting for per-layer costs and transformations between neighboring strategies.The system compares PP degrees and batch sizes; its search complexity is O(L E |S|).Source: §3.3 · p. 5

3.4 Cost Estimation

• Galvatron estimates memory from tensor shape and type, computation from profiled per-sample layer time, and communication from transfer size and bandwidth.For overlapping backward computation and communication, its estimator accounts for GPU resource contention, which slowed both by 1.3× in evaluations.Source: §3.4 · pp. 5–6

4 IMPLEMENTATION

• Users can integrate Galvatron into PyTorch training programs with a few lines of modifications.It is an automatic Transformer training framework and part of the Hetu distributed DL system.Source: §4 · p. 6
• A pre-created global communication-group pool avoids expensive NCCL group-construction overheads for generated plans.The pool contains all GPU groups that might be used.Source: §4 · p. 6
• Slice-Gather places activations on devices as required when neighboring layers use different strategies, recognizing cases needing no communication.A transition from 4-way TP to 4-way DP is given as a zero-communication case.Source: §4 · p. 6

5.1 Experimental Setups

• Experiments compare Galvatron with pure and hybrid parallelism methods across BERT, T5, ViT, and Swin Transformer, using Wikipedia and ImageNet-1K on eight RTX TITAN GPUs.The setup uses Adam; results are averaged over 100 iterations.Source: §5.1 · pp. 6–7

5.2 End-to-End Comparison

• Across model scales and memory budgets, Galvatron outperforms all baselines, reaching up to 338% higher throughput on ViT and up to 55% over hybrid strategies.For the other three models, reported maxima are 200%-334% over single strategies and 28%-52% over hybrid strategies. Model preferences vary, and some DP+TP or DP+PP cases underperform pure SDP.Source: §5.2 · p. 7

5.3 Estimation Performance

• With overlapping slowdown considered, average prediction error is less than 5%; ignoring it yields more than 15%.The comparison covers all experimental models.Source: Figure 3 · p. 7, §5.3 · p. 7

5.4 Optimization Efficiency

• Search time grows linearly with layer count and memory limit, reaching only hundreds of seconds; Galvatron searches 22 candidates versus four each for DP+TP and DP+PP.The Galvatron search takes longer than those limited-strategy searches in the stated comparison.Source: Figure 4 · p. 7, §5.4 · p. 7

5.5 Optimal Parallelism Plan

• For BERT-Huge-32 at 8 GB, Galvatron combines all four basic parallelisms and achieves 75% acceleration; at 12 GB, its plan mixes TP+DP and TP+SDP.For Swin Transformer, shallower layers prefer data parallelism while deeper layers prefer tensor parallelism; Figure 5 encodes parameter size as rectangle height and activation size as width.Source: §5.5 · pp. 7–8, Figure 5 · p. 8

5.6 Scalability Study

• At 16 GPUs, ViT-Huge-32 with 16 GB reaches 131.15 versus 63.25 samples per second; at 64 GPUs, Galvatron still beats baselines on 10-billion-parameter models.Search costs are 2.2× and 9.2× at 16 and 64 GPUs, respectively, compared with 8 GPUs.Source: Table 3 · p. 8, §5.6 · p. 8, Table 4 · p. 8

• Galvatron significantly outperforms state-of-the-art baselines on Transformer training throughput.Source: §6 · p. 8
• Its search-space decomposition and exploration algorithm supports automatic parallel Transformer training across multiple GPUs.Source: §6 · p. 8
• The authors identify expensive training costs as a reason large-scale Transformer training is increasingly important.Source: §6 · p. 8