Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models
Abstract
We present Branch-Train-Merge (BTM), a communication-efﬁcient algorithm for training of language models (LMs).
BTM learns a set of independent EXPERT
LMs (ELMs), each specialized to a different domain, such as scientiﬁc or legal text.
New ELMs are learned by branching from (mixtures of) ELMs in the current set, further training on new domains, and then merging the resulting models back into the set for future use.
These ELMs can be ensembled or averaged at inference time.
Experiments show that BTM improves in- and out-of-domain perplexities as compared to compute-matched GPT-style transformer LMs.
Our results suggest that extreme parallelism could be used to efﬁciently scale LMs in future work.
Introduction
Training and inference in language models (LMs) typically require access to supercomputers that can achieve the massive multi-node synchronization required to compute model activations and gradients (Brown et al., 2020; Fedus et al., 2022; Zhang et al., 2022).
We develop a new class of LMs that is instead embarrassingly parallel: different parts of the model are independently trained on different subsets of the data, with no need for multi-node training or inference (Figure 2).
Our new ELM FOREST 1 model consists of a set of EXPERT LMs (ELMs), independently functional
LMs specialized to a domain in the training corpus, e g., scientiﬁc or legal text, with no shared parameters , which can be ensembled or parameter averaged to collapse back to a single LM at inference time.
Our Branch-Train-Merge (BTM) algorithm learns an ELM FOREST by repeatedly adding new ELMs.
Each new ELM is ﬁrst branched by initializing a new LM with an average of the parameters of relevant LMs in the current set, then further trained on new domains, and ﬁnally merged into the ELM FOREST (Figure 3).
The ELM FOREST is initalized with a single LM, trained on heterogeneous data to establish strong shared representations for future domain specialization.
ELM FOREST s trained with BTM outperform GPT-style transformer LMs and a domain-specialized mixture-of-experts baseline (Gururangan et al. 2022) across a range of computational budgets.
We release code and models.
Methods
We deﬁne an ELM FOREST to be a set of EXPERT LMs (ELMs), each independently trained to specialize to a different subset of a corpus.
ELMs are inspired by the experts in earlier MoE models (Jacobs et al., 1991), but we deﬁne ours to be domain specialists and specialize the full LM.
Following
Gururangan et al.
2022, we deﬁne domains by provenance (the source of the document, e.g., legal document, medical research paper).
3 ELMs remain independent throughout training and inference.
ELM FOREST s support two inference modes.
We perform output ensembling over the output probabilities of multiple ELMs.
Alternatively, we use parameter averaging (Izmailov et al., 2018; Wortsman et al., 2022; Matena & Raffel, 2021) to collapse the ELM FOREST into a single LM, keeping inference cost constant as ELMs are added to the set.
We weight both inference operations with a domain posterior , which estimates the relevance of each expert to the evaluation domain.
2. B RANCH -T RAIN -M ERGE (BTM)
B RANCH -T RAIN -M ERGE training of ELM FOREST models is incremental and embarrassingly parallel; EXPERT LMs are trained fully independently, starting from a seed LM (Appendix Figuresand 3).
Each BTM iteration begins with an existing ELM FOREST E
= { θ i } ki =1 .
Each ELM θ i represents a corresponding domain d i in the dataset of k domains D E
= { d i } k i =1 modeled by E .
We ﬁrst describe the inductive case of k > 0 , then describe how to train the initial model θ 0 .
Step 1 (Branch):
Given some vector of weights w = { w 1 , w 2 , ..., w k } over the existing experts θ 1 , θ 2 , ..., θ k , we initialize the new expert with the weighted parameter average θ k +1 ← (cid:80) ki =
w i θ i .
Step 2 (Train):
We train the new ELM θ k +1 on data domain d k +1 with the log likelihood objective.
None of the existing ELMs in E are involved in the training of the new ELM.
We also refer to this step as branched training to distinguish it from other training regimens.
Step 3 (Merge): We merge the new ELM θ k +1 into E to create an updated set: E (cid:48) =
E ∪ { θ k +1 } .
Step 0 (Initialization):
In the ﬁrst iteration of BTM, E = ∅ ; we have no ELMs in the set to branch from.
Instead of initializing the ﬁrst ELMs randomly, we perform a seed phase , training a seed LM θ seed on some data corpus d seed to initialize the ﬁrst batch of ELMs.
Experiments and Results 3.1 Experimental Setup
Data We use data from Gururangan et al.
(2022), which consists of 8 diverse training andevaluation (primarily English-language) domains.
Details are in Appendix Table 7.
Model hyperparameters The model architecture is a randomly-initialized LM with the GPT- (Brown et al., 2020) architecture implemented in Fairseq (Ott et al., 2019).
We use 125M (small), 350M (medium), 750M (large), 1.3B (xl) parameter models.
Following Brown et al.
2020, we use the
GPT-2 (Radford et al., 2019) vocabulary of 50,264 BPE types, and train with 1,024-token sequences, across document boundaries.
We prepend a beginning-of-document token to each document.
Compared Models For our ELM FOREST s, we ﬁrst conduct a seed phase to initialize the ensemble with LM parameters, then conduct branched training on the ELMs (§2.1), all initialized with the seed LM.
Our baselines are (1) a transformer-LM , implemented with distributed data parallelism (Li, 2021) 5 ; (2) DEM IX (Gururangan et al., 2022), where feedforward layers in the transformer are trained to specialize as domain experts.
These models are compute-matched.
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] Training settings To disentangle variations in GPU speed, we use number of updates as our computational budget: 80k, 32k, 24k, and 12k updates on 16, 32, 64, and 128 GPUs in parallel for the 125M, 350M, 750M, 1.3B parameter TRANSFORMER -LM and DEM IX baselines, respectively.
We use the same GPU counts for the seed phase in the ELM FOREST .
For branched training, we divide these GPUs equally among the ELMs.
For all models, we ﬁx the learning rate at 0.0005 with a polynomial (linear) decay learning rate schedule and 8% warmup, which we found to be optimal for most settings after a large grid search.
We use a batch size of 16 for each GPU, with gradient accumulation of 32 steps, and train with fp16 .
We train on NVIDIA V100 32GB GPUs.
3.2 Results
Ensembling results are shown in Table 1.
At these model scales, ELM FOREST ensembles outperform both the sparsely trained DEM IX LM and the densely trained TRANSFORMER -LM baselines.
While ELM FOREST substantially improves performance at lower training cost relative to the
TRANSFORMER -LM, it comes at the price of a larger model size and higher associated inference costs when ensembling.
Thus, we explore parameter averaging to combine experts for improved generalization with no additional inference costs relative to the TRANSFORMER -LM baseline.
Given some weight vector w over k ELMs { θ i , ..., θ k } , we deﬁne a single model such that its parameters are a weighted average of the ELM parameters, according to w : θ = (cid:80) k i =
w i θ i .
For w , we consider:
Uniform:
We
set
w
to be a uniform; i.e., 1 k .
This setting disregards the relevance of each ELM to the target domain, assuming all ELMs should contribute equally to the average.
Argmax : We set w to be an indicator vector corresponding to the maximum probability in the
domain posterior, thus activating only the estimated best-performing ELM.
Posterior: We set w to be the domain posterior, computed on the validation set.
Results on the evaluation domains are in Table 2.
6 Using uniform weights underperforms all baselines, even TRANSFORMER -LMs, highlighting the importance of domain relevance in output ensembling and parameter averaging ELMs.
Using the argmax ELM outperforms uniform averaging for small models, but not larger models.
Weighting the average with the domain posterior outperforms all other techniques, and consistently improves performance over TRANSFORMER -LMs at no additional inference cost.
Though parameter averaging does not reach the performance of output ensembling, the lower inference costs and simplicity of deployment may make averaging the preferred inference technique for resource-constrained applications.
In Figure 1, we ﬁx the parameter average to use the domain posterior weights and vary the portion of the compute budget dedicated to the seed phase and observe the effects on performance.
We observe that parameter averaging performance on training domains is relatively robust to seed training.
On evaluation domains, however, the smallest scale ELM FOREST does not achieve optimal performance until about 60% or more updates are dedicated to seed training.
This explains the poor performance of the 125M parameter scale ELM FOREST average on evaluation domains in Table 2.
Overall, results suggest a strong effect of the seed phase on the viability of ELM FOREST averaging.
ELM FOREST averaging does not work with ELMs trained from random initialization (i.e., with no seed phase).
[Figure]
Related Work
Sparsely activated language models have been considered in a few forms (Evci et al., 2020; Mostafa &
Wang, 2019; Dettmers & Zettlemoyer, 2019), but most related to this work is the Mixture-of-Experts (MoE) model (Jacobs et al., 1991;
Lepikhin et al., 2021; Fedus et al., 2022; Lewis et al., 2021; Roller et al., 2021).
Of this line of work, ours is most closely related to DEMix layers Gururangan et al. (2022), which replace transformer feedforward layers as domain experts.
Ensemble methods are classic techniques in machine learning (Breiman, 1996; Freund, 1995; Wolpert, 1992).
Similar to our work, recent work has considered growing ensembles, in which new models are trained sequentially on streaming data Caccia et al.
Our averaging mechanism is inspired by the model merging techniques in the vision and NLP literature (Wortsman et al., 2022; Izmailov et al., 2018; Matena & Raffel, 2021).
Our posterior weighted average is highly related to Bayesian model averaging techniques used in classic ensembling methods (Fragoso et al., 2018).
Model averaging has also been explored for federated learning (McMahan et al., 2017).
Conclusion
We introduce BTM, a new algorithm to train an ELM FOREST , which contains many EXPERT LMs.
Our experiments show that ELM FOREST s trained with BTM outperform compute-matched baselines, when conducting inference through output ensembling or parameter averaging.
These results provide compelling evidence for the promise of scaling LMs with many smaller, independently trained ELMs. Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] 2. If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results? [N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Anonymized for review (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] When possible given our compute budget (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] In the appendices and the README of our code repo (currently anonymized) (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] Anonymized (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[No] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[Yes] Brieﬂy – we describe our deidentiﬁcation schema in the appendices 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Additional method details A.1 Finding the domain posterior for inference
Consider the probabilistic view of language modeling, where we estimate p ( X t | x <t ) .
We introduce a domain variable D , alongside each sequence.
Then the next-step conditional distribution on the history x <t is:
[Equation]
We estimate a domain posterior , or a probability of a sequence belonging to each of the k domains using Bayes’ rule:
[Equation]
ELMs are used to compute the likelihood over contexts given a domain label.
To compute the cached prior, we maintain an exponential moving average of posterior probabilities over domains, updated only at the end of each sequence block: p ( D = j )
= (cid:80) Ni = 1 λ i · p ( D = j | x ( i ) < T
[Table] [Table]
Output ensembling naively requires a forward pass through all ELMs in the ELM FOREST , but we observe in practice that the domain posterior is sparse, which that top- k selection of EXPERT
LMs can reduce inference time costs with negligible effects on performance
[Table]
B Efﬁciency Comparison Results
Training ELM FOREST s requires less inter-GPU communication than TRANSFORMER -LM or DEM IX models, since no synchronization occurs between GPUs assigned to different ELMs.
This results in higher updates per second and therefore shorter train times (Table 3).
Additionally, the embarrassingly parallel branched training provides ﬂexibility in resource consumption; GPUs dedicated to different
ELMs may be online at different times, and ELMs may even be trained serially on the same GPUs.
Speciﬁcally, none of our branched training required more than 16 GPUs simultaneously, while our TRANSFORMER -LM training experiments consumed 128 GPUs simultaneously.
Empirically,
ELM FOREST training jobs were scheduled and run more quickly, and with less preemption, than the
TRANSFORMER -LM and DEM IX training jobs at the same overall budget.
C Analysis
In §3, we largely ﬁx the training setup to conduct a controlled comparison of BTM to baseline methods.
We now analyze the importance of various training and inference decisions on language modeling performance.
C. ELM FOREST outperforms parameter-matched ensembles
We compare our method to other LM ensembles to study the effect of increased parameter counts:
Random Ensemble (seed init)
A set of LMs trained on random data splits, to assess the importance of ELM domain specialization.
We pool the training and development sets of the 8 train domains, divide into 8 random splits, then conduct BTM on those splits, with 50% seed training.
[Figure]
ELM FOREST (random init)
An ELM FOREST trained with BTM where all ELMs are randomly initialized, to assess the effect of seed training.
This is equivalent to setting the seed training compute budget to zero updates.
We ﬁx the random initialization across models.
ELM FOREST (seed init)
The ELM FOREST setting of §3.
We conduct BTM on the 8 train domains, and dedicate 50% of the updates in the budget to seed and to branched ELM training.
Results with the largest models are in Table 4.
7 ELM FOREST (random init) nearly matches ELMFOREST on training domains but performs poorly on evaluation domains.
The random ensemble is consistently worse than both variants of ELM FOREST , showing that the performance improvement is not only due to ensembling or increased total model size.
C. ELM FOREST performance is robust to seed LM training compute allocation
The ELM FOREST (random init), which has no seed training, underperforms ELM FOREST (LM init) in §C.1, indicating that seed training is essential.
On the other hand, TRANSFORMER -LM, equivalent to 100% seed training, also underperforms ELM FOREST (LM init) in §3, which suggests the importance of branched ELM training.
We now study the changes to performance when we vary the portion of the compute budget dedicated to seed training.
We control for the total compute budget (across seed and branched training).
Our results, in Figure 4, show that the optimal amount of seed training is about 40–60% of the total budget.
At both ends of the full range, performance deteriorates, approaching the ELM FOREST (random init) and TRANSFORMER -LM performance (at 0% and 100% seed training, respectively).
However, as little as 10% of seed training can be performed to result in strong gains over the
ELM FOREST (random init) and TRANSFORMER -LM.
This suggests that the majority of BTM training may focus on branched training to dramatically reduced computational costs (§B).
The optimal share of compute to use towards each training phase likely depends on many factors, including the total compute budget.
We leave more thorough study of this trend to future work.
[Table]
C. ELM FOREST performance is robust to the choice of seed training corpus
We compare the effects of using different training corpora for seed training in BTM.
Here, we ﬁx the compute budget allocations studied in §C.2 so that 50% of updates are allocated to seed training and 50% to branched training.
As seen in Table 5, our experiments using the most diverse corpora for seed training resulted in the best performance, but even seed training on only JavaScript code yielded better results than the compute-matched TRANSFORMER -LM baseline, and far better than the ELM FOREST (random init) models in Table 1, which use identical random initialization.
This suggests that initializing ELMs with parameters of any model checkpoint is critical.
C.4 Limitations
The deﬁnition of a domain The nature of domains in NLP is a matter of active research.
Textual domains reﬂect language variation that stems from factors such as vocabulary differences (Blitzer et al., 2006), sociolinguistic (Biber, 1988) or demographic (Rickford, 1985; Blodgett et al., 2016) variables, community membership (Lucy & Bamman, 2021), end-tasks (Gururangan et al., 2020), or temporal shifts (Lazaridou et al., 2021;
Luu et al., 2021).
In this work, we follow Gururangan et al. and deﬁne domains by provenance , or the source of the document.
Provenance labels yield simple and interpretable segmentations of a corpus, which are useful for identifying ELMs in our experiments.
However, other methods for discovering domains, including unsupervised techniques (Aharoni & Goldberg, 2020; Chronopoulou et al., 2022), may yield better expert assignments.
We leave experimentation with other deﬁnitions of domain for future work.
Domain posterior data requirement To calculate the domain posteriors used for our ensembling and parameter averaging weights, we assume access to a small additional sample of data to train the vector w .
While it is easy to imagine that extra data may be available for most applications to estimate the posterior, future work may explore the possibility of eliminating this requirement.
Other distributed training baselines Our TRANSFORMER -LM baseline is implemented with distributed data-parallel.
Model-parallel, fully sharded data-parallel, and other distributed training strategies (Artetxe et al., 2021) confer different scaling patterns that may change the conclusions that we report in this work.
However, we expect that BTM will provide strong efﬁciency gains against these alternatives.
Harms of language models BTM results in an LM whose test time behaviors can be controlled with much stronger guarantees after training due to the isolation of domains in ELMs.
However,
ELM FOREST s exposed to large datasets scraped from the Internet may contain toxic language (e.g., hatespeech) that are difﬁcult to identify with coarse provenance domain labels, and nevertheless result in harmful output from the ELMs (Gehman et al., 2020).
Future work may explore recipes for training and deploying ELM FOREST s to better support user safety.
[Table]
https://regex101.com/r/ZqsF9x/ <
[Table]
https://regex101.com/r/0tQ6EN/ <DART>
[Table]
https://regex101.com/r/GZl5EZ/ <FB_USERID
[Table]
https://regex
[Table]
com/r/YrDpPD/ <PHONE
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
125M
[Table]
T
[Table]
T
[Table]
DEM IX
ELM FOREST
R ANDOM
ELM FOREST U
[Table] [Table]
E
[Table]
125M 512M 1B 1B 1B
[Table]
[Table]
ELM
[Table]
ELM
[Table] [Table]
[Table] [Table]