解决大模型训练中4D并行的负载不均问题，提升训练效率。
WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training
- 通过可变长度文档打包平衡微批次计算与通信负载。
- 采用细粒度文档分片策略，确保上下文并行组内各节点负载一致。
- 在不同规模模型上平均提速23%，适合大规模语言模型训练场景。
本文提出WLB-LLM，一种面向大规模语言模型训练的负载均衡4D并行方法。我们深入分析了4D并行训练中的负载不均问题，发现其主要源于流水线并行和上下文并行两个层面。为解决该问题，针对流水线并行，WLB-LLM引入一种基于负载感知的可变长度文档打包方法，以均衡微批次间的计算与通信开销；针对上下文并行，提出一种细粒度的按文档分片策略，确保同一上下文并行组内的每个工作节点承担相同的工作量。在多种模型规模下的综合实验表明，WLB-LLM显著缓解了4D并行训练过程中的负载不均现象，在内部大模型训练框架中实现了平均1.23倍的加速效果。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。