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

Learn about WLB-LLM, a novel workload-balanced 4D parallelism approach for large language model training in this 16-minute conference presentation from OSDI '25. Explore the comprehensive analysis of workload imbalance issues in LLM training, focusing on two primary sources of imbalance at the pipeline parallelism and context parallelism levels. Discover how WLB-LLM addresses these challenges through a workload-aware variable-length document packing method that balances computation and communication workload across micro-batches at the pipeline parallelism level. Examine the innovative fine-grained per-document sharding strategy introduced at the context parallelism level, which ensures identical workload distribution among workers within a context parallelism group. Review comprehensive experimental results across different model scales that demonstrate WLB-LLM's effectiveness in mitigating workload imbalance during 4D parallelism LLM training, achieving an average speedup of 1.23× when integrated into internal LLM training frameworks.