# Improving Automatic Parallel Training via Balanced Memory Workload Optimization

arXiv: https://arxiv.org/abs/2307.02031 · HTML: https://mlsanity.com/paper/2307.02031/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 56 references that are in the map.
- First posted on arXiv 5 Jul 2023.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [BERT4Rec: Sequential Recommendation with Bidirectional Encoder Representations from Transformer](https://mlsanity.com/paper/1904.06690/) (2019)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Do Transformers Really Perform Bad for Graph Representation?](https://mlsanity.com/paper/2106.05234/) (2021)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [PipeMare: Asynchronous Pipeline Parallel DNN Training](https://mlsanity.com/paper/1910.05124/) (2019)
- [Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelism](https://mlsanity.com/paper/2211.13878/) (2022)
- [Supporting Very Large Models using Automatic Dataflow Graph Partitioning](https://mlsanity.com/paper/1807.08887/) (2018)
- [HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System](https://mlsanity.com/paper/2203.14685/) (2022)
- [Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training](https://mlsanity.com/paper/2004.13336/) (2020)
- [HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework](https://mlsanity.com/paper/2112.07221/) (2021)
- [Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms](https://mlsanity.com/paper/2007.00156/) (2020)
- [Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks](https://mlsanity.com/paper/1802.04924/) (2018)

## Built on by (most cited first, 14 of 14)

- [ThunderServe: High-performance and Cost-efficient LLM Serving in Cloud Environments](https://mlsanity.com/paper/2502.09334/) (2025)
- [FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism](https://mlsanity.com/paper/2412.01523/) (2024)
- [Parallax: Efficient LLM Inference Service over Decentralized Environment](https://mlsanity.com/paper/2509.26182/) (2025)
- [Spindle: Efficient Distributed Training of Multi-Task Large Models via Wavefront Scheduling](https://mlsanity.com/paper/2409.03365/) (2024)
- [Apt-Serve: Adaptive Request Scheduling on Hybrid Cache for Scalable LLM Inference Serving](https://mlsanity.com/paper/2504.07494/) (2025)
- [BOute: Cost-Efficient LLM Serving with Heterogeneous LLMs and GPUs via Multi-Objective Bayesian Optimization](https://mlsanity.com/paper/2602.10729/) (2026)
- [LAER-MoE: Load-Adaptive Expert Re-layout for Efficient Mixture-of-Experts Training](https://mlsanity.com/paper/2602.11686/) (2026)
- [MEMO: Fine-grained Tensor Management For Ultra-long Context LLM Training](https://mlsanity.com/paper/2407.12117/) (2024)
- [Cascadia: An Efficient Cascade Serving System for Large Language Models](https://mlsanity.com/paper/2506.04203/) (2025)
- [Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment](https://mlsanity.com/paper/2412.07894/) (2024)
- [OServe: Accelerating LLM Serving via Spatial-Temporal Workload Orchestration](https://mlsanity.com/paper/2602.12151/) (2026)
- [Autopoiesis: A Self-Evolving System Paradigm for LLM Serving Under Runtime Dynamics](https://mlsanity.com/paper/2604.07144/) (2026)
- [HexAGenT: Efficient Agentic LLM Serving via Workflow- and Heterogeneity-Aware Scheduling](https://mlsanity.com/paper/2605.16637/) (2026)
- [Beyond the Trial-and-Error Loop: Hybrid Projection and Automated Tuning for Distributed Training](https://mlsanity.com/paper/2610.19423/) (2026)

## Abstract (arXiv)

Transformer models have emerged as the leading approach for achieving state-of-the-art performance across various application domains, serving as the foundation for advanced large-scale deep learning (DL) models. However, efficiently training these models across multiple GPUs remains a complex challenge due to the abundance of parallelism options. Existing DL systems either require manual efforts to design distributed training plans or limit parallelism combinations to a constrained search space. In this paper, we present Galvatron-BMW, a novel system framework that integrates multiple prevalent parallelism dimensions and automatically identifies the most efficient hybrid parallelism strategy. To effectively navigate this vast search space, we employ a decision tree approach for decomposition and pruning based on intuitive insights. We further utilize a dynamic programming search algorithm to derive the optimal plan. Moreover, to improve resource utilization and enhance system efficiency, we propose a bi-objective optimization workflow that focuses on workload balance. Our evaluations on different Transformer models demonstrate the capabilities of Galvatron-BMW in automating distributed training under varying GPU memory constraints. Across all tested scenarios, Galvatron-BMW consistently achieves superior system throughput, surpassing previous approaches that rely on limited parallelism strategies.