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

TL;DR

This paper introduces DynMo, an autonomous load balancing method for dynamic large language models that reduces workload imbalance and improves training efficiency across distributed systems.

Contribution

DynMo provides a provably optimal load balancing solution that adaptively equalizes compute loads and consolidates computation, enhancing the scalability of dynamic LLM training.

Findings

Achieves up to 4.52x speedup in training with dynamic LLM techniques.

Effectively balances workload across workers in distributed training.

Supports both multi-GPU and multi-node GPU clusters.

Abstract

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.