Branch-Train-MiX : Intégration d'experts LLM dans un modèle de type Mixture-of-Experts LLM

Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

March 12, 2024

Auteurs: Sainbayar Sukhbaatar, Olga Golovneva, Vasu Sharma, Hu Xu, Xi Victoria Lin, Baptiste Rozière, Jacob Kahn, Daniel Li, Wen-tau Yih, Jason Weston, Xian Li

cs.AI

Résumé

Nous étudions des méthodes efficaces pour entraîner des modèles de langage de grande taille (LLMs) afin qu'ils acquièrent des compétences dans plusieurs domaines spécialisés, tels que le codage, le raisonnement mathématique et les connaissances générales. Notre méthode, appelée Branch-Train-MiX (BTX), commence avec un modèle de base, qui est divisé pour entraîner des experts de manière embarrassément parallèle, avec un débit élevé et des coûts de communication réduits. Une fois que les experts individuels sont entraînés de manière asynchrone, BTX regroupe leurs paramètres de propagation avant en tant qu'experts dans des couches de Mélange d'Experts (MoE) et moyenne les paramètres restants, suivis d'une étape de réglage fin MoE pour apprendre le routage au niveau des tokens. BTX généralise deux cas particuliers : la méthode Branch-Train-Merge, qui ne comprend pas l'étape de réglage fin MoE pour apprendre le routage, et le recyclage parcimonieux, qui omet l'étape d'entraînement asynchrone des experts. Par rapport aux approches alternatives, BTX offre le meilleur compromis entre précision et efficacité.

English

We investigate efficient methods for training Large Language Models (LLMs) to possess capabilities in multiple specialized domains, such as coding, math reasoning and world knowledge. Our method, named Branch-Train-MiX (BTX), starts from a seed model, which is branched to train experts in embarrassingly parallel fashion with high throughput and reduced communication cost. After individual experts are asynchronously trained, BTX brings together their feedforward parameters as experts in Mixture-of-Expert (MoE) layers and averages the remaining parameters, followed by an MoE-finetuning stage to learn token-level routing. 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. Compared to alternative approaches, BTX achieves the best accuracy-efficiency tradeoff.