Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

• URL: http://arxiv.org/abs/2501.11873v2
• Date: Tue, 04 Feb 2025 02:07:37 GMT
• Title: Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models
• Authors: Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, Junyang Lin,
• Abstract summary: This paper revisits the implementation of $textbfL$oad-$textbfb$alancing $textbfL$oss (LBL) when training Mixture-of-Experts (MoEs) models.
• Score: 77.79855507792564
• License:
• Abstract: This paper revisits the implementation of $\textbf{L}$oad-$\textbf{b}$alancing $\textbf{L}$oss (LBL) when training Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as $N_E \sum_{i=1}^{N_E} f_i p_i$, where $N_E$ is the total number of experts, $f_i$ represents the frequency of expert $i$ being selected, and $p_i$ denotes the average gating score of the expert $i$. Existing MoE training frameworks usually employ the parallel training strategy so that $f_i$ and the LBL are calculated within a $\textbf{micro-batch}$ and then averaged across parallel groups. In essence, a micro-batch for training billion-scale LLMs normally contains very few sequences. So, the micro-batch LBL is almost at the sequence level, and the router is pushed to distribute the token evenly within each sequence. Under this strict constraint, even tokens from a domain-specific sequence ($\textit{e.g.}$, code) are uniformly routed to all experts, thereby inhibiting expert specialization. In this work, we propose calculating LBL using a $\textbf{global-batch}$ to loose this constraint. Because a global-batch contains much more diverse sequences than a micro-batch, which will encourage load balance at the corpus level. Specifically, we introduce an extra communication step to synchronize $f_i$ across micro-batches and then use it to calculate the LBL. Through experiments on training MoEs-based LLMs (up to $\textbf{42.8B}$ total parameters and $\textbf{400B}$ tokens), we surprisingly find that the global-batch LBL strategy yields excellent performance gains in both pre-training perplexity and downstream tasks. Our analysis reveals that the global-batch LBL also greatly improves the domain specialization of MoE experts.

Related papers

• Towards Efficient Automatic Self-Pruning of Large Language Models [55.90119819642064]
  Post-training structured pruning is a promising solution that prunes Large Language Models without the need for retraining. We argue that the key to mitigating this issue lies in accurately determining the pruning rate for each layer. We introduce $textbfSelf-Pruner$ an end-to-end automatic self-pruning framework for LLMs, which efficiently search layer-wise pruning rates.
  arXiv  Detail & Related papers  (2025-02-20T09:59:50Z)
• Control LLM: Controlled Evolution for Intelligence Retention in LLM [4.67235851066221]
  We propose textbfControl LLM, a novel approach that leverages parallel pre-trained and expanded transformer blocks. Experiments demonstrate the effectiveness of Control LLM in both Continuous Pre-training (CPT) and Continuous Supervised Fine-Tuning (CSFT) It surpasses existing methods and achieves SOTA among open-source models tuned from the same base model, using substantially less data and compute.
  arXiv  Detail & Related papers  (2025-01-19T08:06:06Z)
• Mean-Field Sampling for Cooperative Multi-Agent Reinforcement Learning [4.899818550820576]
  We propose a new algorithm for multi-agent reinforcement learning. We show that this learned policy converges to the optimal policy on the order of $tildeO (1/sqrtk)$ as the number of subsampled agents increases.
  arXiv  Detail & Related papers  (2024-12-01T03:45:17Z)
• FLARE: Faithful Logic-Aided Reasoning and Exploration [50.9814063216852]
  We introduce a novel approach for traversing the problem space using task decompositions. We use the Large Language Models to plan a solution, soft-formalise the query into facts and predicates using a logic programming code. Our method allows us to compute the faithfulness of the reasoning process w.r.t. the generated code and analyse the steps of the multi-hop search without relying on external solvers.
  arXiv  Detail & Related papers  (2024-10-14T19:39:11Z)
• Stacking Your Transformers: A Closer Look at Model Growth for Efficient LLM Pre-Training [42.89066583603415]
  This work identifies three critical $textitO$bstacles: lack of comprehensive evaluation, ($textitO$2) untested viability for scaling, and ($textitO$3) lack of empirical guidelines. We show that a depthwise stacking operator, called $G_textstack$, exhibits remarkable acceleration in training, leading to decreased loss and improved overall performance.
  arXiv  Detail & Related papers  (2024-05-24T08:00:00Z)
• Federated Combinatorial Multi-Agent Multi-Armed Bandits [79.1700188160944]
  This paper introduces a federated learning framework tailored for online optimization with bandit. In this setting, agents subsets of arms, observe noisy rewards for these subsets without accessing individual arm information, and can cooperate and share information at specific intervals.
  arXiv  Detail & Related papers  (2024-05-09T17:40:09Z)
• Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM [81.18305296110853]
  We investigate efficient methods for training Large Language Models (LLMs) to possess capabilities in multiple specialized domains. Our method, named Branch-Train-MiX (BTX), starts from a seed model, which is branched to train experts in embarrassingly parallel fashion. BTX generalizes two special cases, the Branch-Train-Merge method, which does not have the MoE finetuning stage to learn routing, and sparse upcycling, which omits the stage of training experts asynchronously.
  arXiv  Detail & Related papers  (2024-03-12T16:54:58Z)
• Provably Efficient High-Dimensional Bandit Learning with Batched Feedbacks [93.00280593719513]
  We study high-dimensional multi-armed contextual bandits with batched feedback where the $T$ steps of online interactions are divided into $L$ batches. In specific, each batch collects data according to a policy that depends on previous batches and the rewards are revealed only at the end of the batch. Our algorithm achieves regret bounds comparable to those in fully sequential setting with only $mathcalO( log T)$ batches.
  arXiv  Detail & Related papers  (2023-11-22T06:06:54Z)
• Sweeping Heterogeneity with Smart MoPs: Mixture of Prompts for LLM Task Adaptation [43.32632163091792]
  Large Language Models (LLMs) have the ability to solve a variety of tasks, such as text summarization and mathematical questions. Due to high computational costs, the current trend is to use prompt instruction tuning to better adjust monolithic, pretrained LLMs for new -- but often individual -- downstream tasks. MoPs can simultaneously mitigate prompt training "interference" in multi-task, multi-source scenarios.
  arXiv  Detail & Related papers  (2023-10-04T14:11:12Z)
• Batch Model Consolidation: A Multi-Task Model Consolidation Framework [14.687385545898776]
  In Continual Learning (CL), a model is required to learn a stream of tasks sequentially without significant performance degradation on previously learned tasks. We propose Batch Model Consolidation ($textbfBMC$) to support more realistic CL under conditions where multiple agents are exposed to a range of tasks. Our method outperforms the next best CL approach by 70% and is the only approach that can maintain performance at the end of 71 tasks.
  arXiv  Detail & Related papers  (2023-05-25T21:33:56Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.