Branch-Train-MiX: Experten-LLMs in einen Mixture-of-Experts-LLM mischen

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

March 12, 2024

Autoren: 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

Zusammenfassung

Wir untersuchen effiziente Methoden zur Schulung von Large Language Models (LLMs), um Fähigkeiten in mehreren spezialisierten Bereichen wie Codierung, mathematischer Logik und Weltwissen zu erlangen. Unsere Methode, namens Branch-Train-MiX (BTX), beginnt mit einem Ausgangsmodell, das verzweigt wird, um Experten in peinlich paralleler Weise mit hoher Durchsatzrate und reduzierten Kommunikationskosten zu schulen. Nachdem die einzelnen Experten asynchron geschult wurden, bringt BTX ihre Feedforward-Parameter als Experten in Mixture-of-Expert (MoE)-Schichten zusammen und mittelt die verbleibenden Parameter, gefolgt von einer MoE-Feinabstimmungsphase, um das Routing auf Token-Ebene zu erlernen. BTX verallgemeinert zwei spezielle Fälle, die Branch-Train-Merge-Methode, die keine MoE-Feinabstimmungsphase zum Erlernen des Routings hat, und das sparse upcycling, das die Phase des asynchronen Trainings von Experten auslässt. Im Vergleich zu alternativen Ansätzen erzielt BTX das beste Verhältnis von Genauigkeit und Effizienz.

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.