Research map / Architectures and efficiency / RNNs and normalization
UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programming
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes UniAP, an automatic parallelism method that jointly optimizes inter-layer (pipeline) and intra-layer (DP, TP, FSDP) strategies by formulating the problem as a mixed integer quadratic program solved with Gurobi.
Problem
Existing automatic parallelism methods optimize inter-layer and intra-layer parallel strategies separately or hierarchically, yielding sub-optimal distributed training strategies, and manual parallelism demands heavy human effort.
Approach
Encode pipeline partitioning, per-layer intra-layer strategy selection, resharding and memory constraints, and order-preserving contiguity jointly as one MIQP over a time cost model, and iterate over pipeline degree and micro-batch count to find the minimal time per iteration.
Key results
- On ENVC, UniAP achieves 4.63 samples/s on Llama-7B versus 1.22 for Galvatron (3.80x throughput) and 11.83x faster strategy optimization.
- Strategy optimization time reduced by up to 107x (vs. >40 min for Alpa on BERT-Huge; 0.37 min for UniAP).
- On ENVD-32, UniAP reaches 3.07 samples/s on Llama-7B versus 2.01 for Megatron, with about 157x faster optimization than Megatron.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVC, Llama-7B | training throughput: 1.22 ± 0.01 -> 4.63 ± 0.007 | Helped |
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVB-8, BERT-Huge | training throughput: 8.95 ± 0.06 -> 10.77 ± 0.13 | Helped |
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVB-8, ViT-Huge | training throughput: 38.74 ± 0.20 -> 45.58 ± 0.54 | Helped |
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVB-8, Swin-Huge | training throughput: 13.90 ± 0.17 -> 19.08 ± 0.10 | Helped |
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVB-8, T5-Large | training throughput: 8.06 ± 0.06 -> 7.98 ± 0.05 | Hurt |
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVA, BERT-Huge | training throughput: 33.46 ± 0.28 -> 33.46 ± 0.28 | No clear effect |
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVA, T5-Large | training throughput: 23.29 ± 0.04 -> 23.29 ± 0.04 | No clear effect |
| Jointly optimize inter- and intra-layer parallelism | UniAP · ENVA, ViT-Huge | training throughput: 109.51 ± 0.07 -> 109.51 ± 0.07 | No clear effect |
| Constrain the strategy space to inter-layer-only strategies (ablation) | UniAP · ENVB, BERT-Huge-32; T5-Large-32; ViT-Huge-32; Swin-Huge-48 | training throughput | Mixed |
| Constrain the strategy space to intra-layer-only strategies (ablation) | UniAP · ENVB, BERT-Huge-32; T5-Large-32; ViT-Huge-32; Swin-Huge-48 | training throughput | Mixed |
Surprising findings
- With Megatron on Llama-7B, 64.1% of candidate parallel strategies fail with OOM, showing manual parallelism is error-prone.
- Galvatron's hierarchical search on ENVC chose a 4-stage PP with 2-way TP, while UniAP's joint search picked 8-stage PP with micro-batch size 1, outperforming it by 3.80x.
Limitations
- Relies on an analytical cost model and profiling, with solver time limits; optimality is relative to the model.
- Uses a GPipe-style pipeline objective; other pipeline schedules require modifying the memory constraint.
- Evaluated on specific clusters and five Transformer-based models only; Alpa and Swin/Llama implementations were unavailable or failed to compile.
Experimental setup
Metrics: training throughput (samples/s) strategy optimization time (min.)
Baselines: Galvatron Alpa Megatron-LM DeepSpeed (ZeRO-3)
Builds on: GPipe Galvatron Alpa Megatron-LM Gurobi
Closest papers by meaning
- Deep Distributed Optimization for Large-Scale Quadratic Programming (ICLR 2025 · similarity 0.86)
- A Quadratic Synchronization Rule for Distributed Deep Learning (ICLR 2024 · similarity 0.84)
- Pfeife: Automatic Pipeline Parallelism for PyTorch (ICML 2025 · similarity 0.84)
- Towards a Better Theoretical Understanding of Independent Subnetwork Training (ICML 2024 · similarity 0.84)
- Memory-Efficient Gradient Unrolling for Large-Scale Bi-level Optimization (NeurIPS 2024 · similarity 0.84)
- Synergistic Tensor and Pipeline Parallelism (NeurIPS 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.