WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _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 |