WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training
这篇分析了工作负载不平衡的来源。输入的sequence是由不同的文本构成的，所以对于同样长度的sequence，文本的长度可以不同。因为attention mask的原因，长文本的尾部token需要大量的计算，这就是导致同长度序列计算开销不同的原因。
这种问题在4D并行中更为明显，现有的分配方式是平均分相同数量的token，没有考虑上面原因导致的计算量异构。
比较简单的一种想法是混合起来一起重构，就像上面右边的图一样。但是这样是跨多个batch在混合重构，一个是效率不太行，另外打乱了原本训练集的随机性，导致收敛出现问题。
