Global-Batch Load Balance: Bijna Gratis Lunch om MoE LLM-training te Verbeteren

De Mixture-of-Experts (MoE)-architectuur is een populaire techniek om modelparameters te schalen, maar bestaande trainingsframeworks zoals Megatron-core gebruiken load balancing-verlies op micro-batchniveau, wat de prestaties kan schaden wanneer micro-batches homogene data bevatten (bijvoorbeeld alleen code). Alibaba Qwen introduceert een global-batch balance loss, die de frequentie van expertselectie synchroniseert over alle parallelle groepen en het verlies wereldwijd berekent. Experimenten met drie MoE-configuraties (3,4B/0,6B geactiveerd, 15B/2,54B geactiveerd, 43B/6,6B geactiveerd) en maximaal 400B tokens tonen aan dat global-batch balance consistent beter presteert dan micro-batch balance op diverse taken. Het leidt ook tot significante domeinspecialisatie van experts, die ontbreekt bij modellen gebalanceerd op micro-batchniveau. Het toevoegen van een klein micro-batch balance loss (gewicht 0,01) bovenop het global loss verbetert de trainingssnelheid zonder de effectiviteit te beïnvloeden. De methode is bijna gratis omdat de synchronisatie alleen een vector van grootte gelijk aan het aantal experts omvat. Het werk is gepubliceerd in arXiv preprint 2501.11873.