# Balanced and Elastic End-to-end Training of Dynamic LLMs

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Serving systems and request scheduling, one of 12 directions in Efficient LLM inference and serving.
- Builds on 33 of its 59 references that are in the map.
- First posted on arXiv 20 May 2025.

## Placed in

- Efficient LLM inference and serving / [Serving systems and request scheduling](https://mlsanity.com/inference/serving-systems-and-request-scheduling/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [To prune, or not to prune: exploring the efficacy of pruning for model compression](https://mlsanity.com/paper/1710.01878/) (2017)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)
- [The State of Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1902.09574/) (2019)
- [Dynamic Neural Networks: A Survey](https://mlsanity.com/paper/2102.04906/) (2021)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [Depth-Adaptive Transformer](https://mlsanity.com/paper/1910.10073/) (2019)
- [Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts](https://mlsanity.com/paper/2408.15664/) (2024)
- [Tutel: Adaptive Mixture-of-Experts at Scale](https://mlsanity.com/paper/2206.03382/) (2022)
- [Compute Trends Across Three Eras of Machine Learning](https://mlsanity.com/paper/2202.05924/) (2022)
- [Sparse GPU Kernels for Deep Learning](https://mlsanity.com/paper/2006.10901/) (2020)
- [Deep Rewiring: Training very sparse deep networks](https://mlsanity.com/paper/1711.05136/) (2017)
- [Learned Token Pruning for Transformers](https://mlsanity.com/paper/2107.00910/) (2021)
- [LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training](https://mlsanity.com/paper/2406.16554/) (2024)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)
- [When BERT Plays the Lottery, All Tickets Are Winning](https://mlsanity.com/paper/2005.00561/) (2020)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines](https://mlsanity.com/paper/2107.06925/) (2021)
- [AutoFreeze: Automatically Freezing Model Blocks to Accelerate Fine-tuning](https://mlsanity.com/paper/2102.01386/) (2021)
- [PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers](https://mlsanity.com/paper/2102.03161/) (2021)
- [Not All GPUs Are Created Equal: Characterizing Variability in Large-Scale, Accelerator-Rich Systems](https://mlsanity.com/paper/2208.11035/) (2022)

## Abstract (arXiv)

To reduce the computational and memory overhead of Large Language Models, various approaches have been proposed. These include a) Mixture of Experts (MoEs), where token routing affects compute balance; b) gradual pruning of model parameters; c) dynamically freezing layers; d) dynamic sparse attention mechanisms; e) early exit of tokens as they pass through model layers; and f) Mixture of Depths (MoDs), where tokens bypass certain blocks. While these approaches are effective in reducing overall computation, they often introduce significant workload imbalance across workers. In many cases, this imbalance is severe enough to render the techniques impractical for large-scale distributed training, limiting their applicability to toy models due to poor efficiency. We propose an autonomous dynamic load balancing solution, DynMo, which provably achieves maximum reduction in workload imbalance and adaptively equalizes compute loads across workers in pipeline-parallel training. In addition, DynMo dynamically consolidates computation onto fewer workers without sacrificing training throughput, allowing idle workers to be released back to the job manager. DynMo supports both single-node multi-GPU systems and multi-node GPU clusters, and can be used in practical deployment. Compared to static distributed training solutions such as Megatron-LM and DeepSpeed, DynMo accelerates the end-to-end training of dynamic GPT models by up to 1.23x for MoEs, 3.18x for parameter pruning, 2.23x for layer freezing, 4.02x for sparse attention, 4.52x for early exit, and 1.17x for MoDs.