WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866910889789095936 |
|---|---|
| author | Wang, Zheng Cai, Anna Xie, Xinfeng Pan, Zaifeng Guan, Yue Chu, Weiwei Wang, Jie Li, Shikai Huang, Jianyu Cai, Chris Hao, Yuchen Ding, Yufei |
| author_facet | Wang, Zheng Cai, Anna Xie, Xinfeng Pan, Zaifeng Guan, Yue Chu, Weiwei Wang, Jie Li, Shikai Huang, Jianyu Cai, Chris Hao, Yuchen Ding, Yufei |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2503_17924 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training Wang, Zheng Cai, Anna Xie, Xinfeng Pan, Zaifeng Guan, Yue Chu, Weiwei Wang, Jie Li, Shikai Huang, Jianyu Cai, Chris Hao, Yuchen Ding, Yufei Distributed, Parallel, and Cluster Computing Artificial Intelligence Machine Learning I.2.11 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. |
| title | WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training |
| topic | Distributed, Parallel, and Cluster Computing Artificial Intelligence Machine Learning I.2.11 |
| url | https://arxiv.org/abs/2503.17924 |