mlsanityData to 14 Oct 2026

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

arXiv 2503.17924 ↗First posted 23 Mar 2025cs.DC cs.AI cs.LGcited by 17

Not on the map. Why?

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.

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

23 of 39 references.

References · 23

  1. 2023GPT-4 Technical Reportcited by 18,396
  2. 2024The Llama 3 Herd of Modelscited by 14,715
  3. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
  4. 2020Scaling Laws for Neural Language Modelscited by 6,586
  5. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  6. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
  7. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
  8. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
  9. 2022CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesiscited by 885
  10. 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 660
  11. 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 632
  12. 2022Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Modelcited by 559
11 more
  1. 2023Ring Attention with Blockwise Transformers for Near-Infinite Contextcited by 457
  2. 2021ZeRO-Offload: Democratizing Billion-Scale Model Trainingcited by 341
  3. 2022Reducing Activation Recomputation in Large Transformer Modelscited by 302
  4. 2023Prompting Large Language Model for Machine Translation: A Case Studycited by 250
  5. 2023DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Modelscited by 245
  6. 2021TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Modelscited by 75
  7. 2023Copiloting the Copilots: Fusing Large Language Models with Completion Engines for Automated Program Repaircited by 55
  8. 2023ZeRO++: Extremely Efficient Collective Communication for Giant Model Trainingcited by 52
  9. 2022ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputscited by 32
  10. 2022Boosting Distributed Training Performance of the Unpadded BERT Modelcited by 15
  11. 2023DynaPipe: Optimizing Multi-task Training through Dynamic Pipelinescited by 12

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 8; 2026: 9.

17 papers.

Sort:

By area:

  1. 2026Scalable Training of Mixture-of-Experts Models with Megatron Corecited by 19
  2. 2025DCP: Addressing Input Dynamism In Long-Context Training via Dynamic Context Parallelismcited by 7
  3. 2025Cornfigurator: Automated Planning for Any-to-Any Multimodal Model Servingcited by 4
  4. 2025WeChat-YATT: A Scalable, Simple, Efficient, and Production Ready Training Librarycited by 2
  5. 2025OctoPipe: Reducing Pipeline Bubbles for Heterogeneous Models via Co-Optimizing Partitioning, Placement, and Schedulingcited by 2
  6. 2026Revisiting Parameter Server in LLM Post-Trainingcited by 2
  7. 2025MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Trainingcited by 1
  8. 2026JigsawRL: Assembling RL Pipelines for Efficient LLM Post-Trainingcited by 1
  9. 2026FlashCP: Load-Balanced Communication-Efficient Context Parallelism for LLM Trainingcited by 1
  10. 2025Concertina: Data-Centric Adaptive Pipeline Parallelism for Efficient Heterogeneous Long-Context LLM Trainingcited by 0
  11. 2025SlimPack: Fine-Grained Asymmetric Packing for Balanced and Efficient Variable-Length LLM Trainingcited by 0
  12. 2025Efficient Long-context Language Model Training by Core Attention Disaggregationcited by 0
5 more
  1. 2026Horizon-LM: A RAM-Centric Architecture for LLM Trainingcited by 0
  2. 2026SparseBalance: Load-Balanced Long Context Training with Dynamic Sparse Attentioncited by 0
  3. 2026ASAP: A Disaggregated and Asynchronous Inference System for MoE Prefillcited by 0
  4. 2026The Cost and Network Limits of Space-Based AI Computecited by 0
  5. 2026The Life of a Token: from Words to Bits on the Wirecited by 0