# WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 39 references that are in the map.
- First posted on arXiv 23 Mar 2025.

## Builds on (most cited first, 23 of 23)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesis](https://mlsanity.com/paper/2203.13474/) (2022)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [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)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [Prompting Large Language Model for Machine Translation: A Case Study](https://mlsanity.com/paper/2301.07069/) (2023)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [Copiloting the Copilots: Fusing Large Language Models with Completion Engines for Automated Program Repair](https://mlsanity.com/paper/2309.00608/) (2023)
- [ZeRO++: Extremely Efficient Collective Communication for Giant Model Training](https://mlsanity.com/paper/2306.10209/) (2023)
- [ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs](https://mlsanity.com/paper/2210.03052/) (2022)
- [Boosting Distributed Training Performance of the Unpadded BERT Model](https://mlsanity.com/paper/2208.08124/) (2022)
- [DynaPipe: Optimizing Multi-task Training through Dynamic Pipelines](https://mlsanity.com/paper/2311.10418/) (2023)

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

- [Scalable Training of Mixture-of-Experts Models with Megatron Core](https://mlsanity.com/paper/2603.07685/) (2026)
- [DCP: Addressing Input Dynamism In Long-Context Training via Dynamic Context Parallelism](https://mlsanity.com/paper/2510.10620/) (2025)
- [Cornfigurator: Automated Planning for Any-to-Any Multimodal Model Serving](https://mlsanity.com/paper/2512.14098/) (2025)
- [WeChat-YATT: A Scalable, Simple, Efficient, and Production Ready Training Library](https://mlsanity.com/paper/2508.07970/) (2025)
- [OctoPipe: Reducing Pipeline Bubbles for Heterogeneous Models via Co-Optimizing Partitioning, Placement, and Scheduling](https://mlsanity.com/paper/2509.23722/) (2025)
- [Revisiting Parameter Server in LLM Post-Training](https://mlsanity.com/paper/2601.19362/) (2026)
- [MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training](https://mlsanity.com/paper/2510.18830/) (2025)
- [JigsawRL: Assembling RL Pipelines for Efficient LLM Post-Training](https://mlsanity.com/paper/2604.23838/) (2026)
- [FlashCP: Load-Balanced Communication-Efficient Context Parallelism for LLM Training](https://mlsanity.com/paper/2606.08476/) (2026)
- [Concertina: Data-Centric Adaptive Pipeline Parallelism for Efficient Heterogeneous Long-Context LLM Training](https://mlsanity.com/paper/2509.21275/) (2025)
- [SlimPack: Fine-Grained Asymmetric Packing for Balanced and Efficient Variable-Length LLM Training](https://mlsanity.com/paper/2509.26246/) (2025)
- [Efficient Long-context Language Model Training by Core Attention Disaggregation](https://mlsanity.com/paper/2510.18121/) (2025)
- [Horizon-LM: A RAM-Centric Architecture for LLM Training](https://mlsanity.com/paper/2602.04816/) (2026)
- [SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attention](https://mlsanity.com/paper/2604.13847/) (2026)
- [ASAP: A Disaggregated and Asynchronous Inference System for MoE Prefill](https://mlsanity.com/paper/2606.22541/) (2026)
- [The Cost and Network Limits of Space-Based AI Compute](https://mlsanity.com/paper/2607.14172/) (2026)
- [The Life of a Token: from Words to Bits on the Wire](https://mlsanity.com/paper/2609.19924/) (2026)

## Abstract (arXiv)

In this work, we present WLB-LLM, a workLoad-balanced 4D parallelism for large language model training. We first thoroughly analyze the workload imbalance issue in LLM training and identify two primary sources of imbalance at the pipeline parallelism and context parallelism levels. Then, to address the imbalance issue, at the pipeline parallelism level, WLB-LLM incorporates a workload-aware variable-length document packing method to balance the computation and communication workload across micro-batches. Additionally, at the context parallelism level, WLB-LLM introduces a novel fine-grained per-document sharding strategy, ensuring each worker within a context parallelism group has an identical workload. Comprehensive experiments under different model scales demonstrate that WLB-LLM significantly mitigates the workload imbalance during 4D parallelism LLM training and achieves an average speedup of 1.23x when applying WLB-LLM in our internal LLM training framework.