We present Branch-Train-Merge (BTM) , a communication-efficient algorithm for embarrassingly parallel training of large language models (LLMs) .
We show it is possible to independently train subparts of a new class of LLMs on different subsets of the data, eliminating the massive multi-node synchronization currently required to train LLMs.BTMlearns a set of independentexpert LMs (ELMs) , each specialized to a different textual domain, such as scientific or legal text. TheseELMs can be added and removed to update data coverage, ensembled to generalize to new domains, or averaged to collapse back to a single LM for efficient inference. NewELMs are learned bybranchingfrom (mixtures of) ELMs in the current set, furthertrainingthe parameters on data for the new domain, and thenmergingthe resulting model back into the set for future use.
Experiments show thatBTMimproves in- and out-of-domain perplexities as compared to GPT-style Transformer LMs, when controlling for training cost. Through extensive analysis, we show that these results are robust to differentELMinitialization schemes, but require expert domain specialization; LM ensembles with random data splits do not perform well. We also present a study of scalingBTMinto a new corpus of 64 domains (192B whitespace-separated tokens in total) ; the resulting LM (22.4B total parameters)  performs as well as a Transformer LM trained with 2.5more compute. These gains grow with the number of domains, suggesting more aggressive parallelism could be used to efficiently train larger models in future work.