Language Adaptation on a Tight Academic Compute Budget: Tokenizer Swapping Works and Pure bfloat16 Is Enough

Authors: konstantin  dobler
gerard  de melo
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-05-31 16:29:46

Abstract: We investigate continued pretraining of LLMs for language adaptation on a tight academic budget: a setting in which only a few GPUs can be used in parallel, for a heavily constrained duration. We focus on adapting Mistral-7B to German or Arabic and evaluate several techniques to improve efficiency and effectiveness in this setting. Our German models adapted on this tight compute budget underperform compared to the base Mistral-7B, while our Arabic models outperform several baselines, showing that for sufficiently well-represented languages, continued pretraining for specialization is not always helpful. Our main findings focus on training precision and tokenizer swapping. Our results show that pure bfloat16 training is a viable alternative to mixed-precision training, while being much faster when only using a few GPUs. Swapping the tokenizer for a specialized one yields more efficient tokenization and is competitive with the original tokenizer, which already contains some German tokens, but did not significantly increase performance for German. Code and model weights are available on GitHub.

Main paper: 
 introduction: 
 Large language models (LLMs) can be incredibly useful tools -if they handle your language well. Unfortunately, for people who speak languages that are not among the small subset of high-resource languages, this is often not the case. One remedy is to adapt an existing LLM to your desired target language via continued pretraining (Chau et al., 2020; Gururangan et al., 2020) . This leverages the enormous amount of compute spent on training the base LLM. Still, as models have gotten bigger, even the compute required for continued pretraining can be too much for many academic labs to handle, requiring the support of large governmental or industrial compute grants.In this work, we use the term "tight academic (compute) budgets", referring to constrained access to a limited number of GPUs (such as only two or four GPUs), constrained GPU memory capacity (e.g., 40GB vs. 80GB Nvidia A100 GPUs), and constrained access durations. 1 1 As academic labs commonly do have access to server-grade rather than consumer-grade GPUs, our definition does encompass the availability of (a few) server-grade GPUs, such as Nvidia A100s, in this setting. We investigate language adaptation of LLMs on such a tight academic budget. We start with the training recipe of the recent LLM adaptation project LeoLM (Plüster et al., 2023) adapting Mistral-7B2 (Jiang et al., 2023a) to German, and modify it for our tight academic compute budget setting. Selecting German as the target language allows us to draw a rough lower bound on the viability of language adaptation on a tight academic budget since Mistral-7B has already seen German during pretraining. For unseen languages, the benefits of language adaptation will be more pronounced. We verify this in an additional "hindsight study" targeting Arabic after our main experiments in German. The focus of our study is not to produce a new state-of-the-art language model in German or Arabic. Rather, we investigate the viability of efficient training techniques to inform future language adaptation attempts on a tight compute budget.We specifically focus on two dimensions of efficiency: tokenizer swapping and training precision. Firstly, we investigate the necessity of the commonly used mixed-precision training (Micikevicius et al., 2017) over training in pure bfloat16 precision. In settings with only two or four GPUs being used in parallel for adapting 7 billion parameter models, mixed-precision bfloat16 training will run out of memory or is only possible if used with inefficient memory-saving techniques like activation checkpointing. This is avoided in pure bfloat16 training by eliminating the need to store full-precision optimizer states and a model weight copy, as required for mixed precision.Secondly, we investigate swapping out Mistral-7B's original tokenizer with a specialized German tokenizer. Many recent language adaptation studies choose to retain the base LLM's original tokenizer (e.g., Kuulmets et al., 2024; Huang et al., 2023; Plüster et al., 2023) . Others opt to extend the original vocabulary rather than replacing it (Zhao et al., 2024b; Fujii et al., 2024; Nguyen et al., 2023) . A more specialized tokenizer has better fertility on our training data, allowing us to train on a greater total number of words (although not tokens) for the same compute and offers a semantically (more) sensible tokenization of text. Swapping instead of extending the tokenizer, resulting in smaller embedding matrices, further boosts training efficiency due to the faster unembedding matrix multiplication and softmax operations and a smaller memory footprint.We summarize our results as follows:• We recommend pure bfloat16 training with small caveats (see Section 4.1) for its large training efficiency gains, especially in tight academic compute budget settings with fewer than eight parallel GPUs, challenging the commonly used mixed-precision bfloat16 paradigm for continued training with extensive analysis and experiments.• We find tokenizer swapping performs on par with keeping the original tokenizer even on a small compute budget, although it does not improve task performance (see Section 4.2), while yielding more efficient tokenization.• Continued pretraining of Mistral-7B on German decreased German task performance, whereas adapting to Arabic led to significant gains. This demonstrates that language adaptation is not always beneficial when the target language is already well-represented (see Section 4.3).We further provide an introduction to mixed-precision LLM training in Section 2. We provide more details on our experimental setup in Section 3 and present the results in Section 4. We discuss related work in Section 5 and the limitations of this study in Section 6. We publish our code and model checkpoints under https://github.com/ konstantinjdobler/tight-budget-llm-adaptation. 
background: precision types: 
 One of our main research questions is whether mixedprecision training is necessary for continued pretraining.8 bits 23 bits ... M M M8 bits 7 bitsFigure 1 : Illustration of the memory layout of float32 and bfloat16, based on García-Nava et al. (2022) .In this section, we provide a short introduction to mixedprecision bfloat16 training. The educated reader is encouraged to skip ahead to Section 3.Nowadays, the most common training setup is bfloat16 mixed-precision training. 3 In mixed-precision training, the forward and backward pass is computed in bfloat16 instead of float32, as this increases computational efficiency.However, bfloat16 has less precision than float32 due to its memory layout with fewer mantissa bits (as illustrated in Figure 1 ). Hence, a master copy of the model weights and the optimizer states are stored in float32 and used during the optimizer step to offset this reduced precision, at the cost of increased memory usage.In settings with constrained access to GPUs, such as our tight academic budget setting, this additional memory cost can force the use of smaller batch sizes and activation checkpointing, which result in significantly lower training efficiency. For example, in our experiments with two 80GB GPUs, pure bfloat16 was 39% faster than mixed-precision bfloat16 training. With access to only a single 80GB GPU, mixed precision proved impossible due to out-of-memory errors. Therefore, we extensively analyze the viability of pure bfloat16 training for language adaptation in the tight academic compute budget setting.The Numerics of bfloat16. Let us take a closer look at how floating point numbers are encoded. The formula for the actual value of a floating point number n is:n = s × 2 e-127 + s × 2 e-127 × 2 -mwhere s is the sign (-1 or 1), e is the number encoded by the exponent bits, and m is the number encoded by the mantissa bits. An intuitive model of this is that we have base numbers represented by the first term s × 2 e-127 that are powers of two over the entire representable range of the datatype (e.g., 2,-2, ... 64,-64, ...) and "fractional" numbers represented by the second term s × 2 e-127 × 2 -m that add a fraction of that power of two until the next power of two is reached.Note that we always have the same amount of fractions between any two powers of two. For bfloat16, we always have 128 = 2 7 fractions, since we have 7 mantissa bits. Therefore, bfloat16 will have very low precision for very large numbers. We have to cover the space between 256 and 512 with the same number of fractions as for the space between 2 and 4. This means that we can only represent 128 floating point numbers between 256 and 512, which only covers every second integer. All other numbers get rounded to the nearest representable one. The maximum ϵ so that x + ϵ still results in x grows as x grows, or -in other words -if x is larger, the minimum representable change of x is also larger.Large weights lead to problems when training in pure bfloat16. When we have some small update u that we want to add/subtract from a weight W , because of the (im-)precision of bfloat16, this update has no effect when u < W/128 (128 = 2 7 for the 7 mantissa bits). This is because u is not large enough to push W to the next fractional value that is representable in bfloat16. So how big of a problem is this in practice? Usually, our network's weights are not that big. If an individual parameter's value is 0.05, then the smallest possible update that does not disappear due to bfloat16 numerics is 0.05/128 ≈ 0.00039. As a side note: the same effect exists for training in float32, but only when u < W/2 23 for 23 mantissa bits in float32. Arguably, this is not a problem in practice.Case study: Mistral-7B. In Figure 2 , we plot the absolute values of Mistral-7B's weights. We see that most weight values are relatively small, which is good, as even very small update values will not disappear in pure bfloat16 training. However, when considering the average values of RMSNorm weights and other weights separately, we observe that the RMSNorm weight values are much larger than other weights. This means that small updates to their weights during training in pure bfloat16 are more likely to disappear. We evaluate and discuss this in Section 4.1. 
experimental setup: 
 We now describe the experimental setup of our main experiments and hindsight study. (Plüster et al., 2023) . The learning rate schedule is modified in our hindsight study, see Section 3.2.We build on the training recipe from LeoLM, which adapts Mistral-7B to German via full fine-tuning of all parameters. We report our hyperparameters in Table 1 . In particular, we use a context length of 4,096 rather than 8,192 tokens to reduce the required GPU memory, as we cannot shard the model across a large number of GPUs. Additionally, we only train for half the number of optimizer steps and use a total batch size of 256 instead of 512 to fit our tight academic compute budget setting. These are our main adjustments to reduce the amount of total compute necessary to complete the training. As a result, our training recipe trains for 8 billion tokens compared to LeoLM's 64 billion tokens, which amounts to 12.5%.The LeoLM recipe is intended to produce a bilingual model by mixing in English data during continued pretraining and using a lower learning rate, which has been shown to reduce catastrophic forgetting but decreases adaptation to the target domain (Ibrahim et al., 2024) . In our work, we instead choose to focus solely on the target language. This allows us to use a higher learning rate, which has also been shown to help adaptation (Ibrahim et al., 2024) . After initial probing experiments4 , we double the learning rate to 4 × 10 -5 .We use the data from the German split of OS-CAR23.01 (Abadji et al., 2022) and filter out all documents with quality warnings. 5 As noted, we do not mix in additional English data or code. For tokenizer swapping, we train a sentencepiece (Kudo & Richardson, 2018 ) BPE tokenizer with byte-fallback on a subset of our training data with a vocabulary size of 32,768 tokens for optimized Nvidia Tensor Core usage. Importantly for training sentencepiece tokenizers on web-crawled data, we use a character coverage of slightly less than 100% to prevent very infrequent characters from being added to the vocabulary (Dobler & de Melo, 2023) .For downstream task evaluation, we rely on a suite of German benchmarks by Plüster (2023) , which were also used for LeoLM. We use their fork of the lm-eval-harness (Gao et al., 2023) to evaluate our checkpoints. The individual datasets used are translated versions of MMLU (Hendrycks et al., 2021) , TruthfulQA (Lin et al., 2021) , ARC (Clark et al., 2018a) , HellaSwag (Zellers et al., 2019) , as well as PAWS-X (Yang et al., 2019) and LAMBADA-OpenAI (Paperno et al., 2016; Radford et al., 2019) .We perform the computation of this work in a real-world tight academic compute budget setting. We use either two or four Nvidia A100 80GB, A100 40GB, H100 80GB, or A6000 48GB GPUs for parallel training with PyTorch FSDP (Zhao et al., 2023) . Depending on the number of available GPUs and training precision, we adjust efficiency settings, such as micro-batch size, activation checkpointing, or different ZeRO (Rajbhandari et al., 2019) stages of Py-Torch FSDP as needed to fit the available GPU memory. We use the FlashAttention self-attention and RMSNorm CUDA kernels (Dao, 2023) .We conduct an additional hindsight study using pure bfloat16 and tokenizer swapping based on the findings from our main experiments for further analysis. We additionally make the following changes to our original training recipe: (i) We use training data from CulturaX (Nguyen et al., 2024) For the Arabic downstream tasks, we evaluate on the OALL (Elfilali et al., 2024) benchmark suite using lighteval (Fourrier et al., 2023) . The individual benchmarks are ACVA (Huang et al., 2023) , AlGhafa (Almazrouei et al., 2023) , and translated Arabic benchmarks from AlGhafa-T based on: MMLU (Koto et al., 2024; Hendrycks et al., 2021) , EXAMS (Hardalov et al., 2020) , ARC-Challenge (Clark et al., 2018b) , ARC-Easy (Clark et al., 2018b) , BOOLQ (Clark et al., 2019) , COPA (Roemmele et al., 2011) , HellaSwag (Zellers et al., 2019) , Open-BookQA (Mihaylov et al., 2018) , PIQA (Bisk et al., 2020) , RACE (Lai et al., 2017) , SciTail (Welbl et al., 2017) , TOXI-GEN (Hartvigsen et al., 2022) . We report a macro average as in OALL (Elfilali et al., 2024) counting the benchmarks in AlGhafa-T individually. We provide further details in Appendix D. 
results & analysis: 
 Now we present and discuss the results of our experiments. We analyze pure vs. mixed-precision bfloat16 in Section 4.1, tokenizer swapping in Section 4.2, and the influence of the target language in Section 4.3. For a single 80GB GPU, mixed-precision bfloat16 training was impossible even when applying all memory saving techniques due to running out-of-memory (OOM). When using only two parallel 80GB GPUs, it became necessary to use a paged AdamW implementation (Dettmers et al., 2023) or to sync gradients at every step during gradient accumulation7 . Otherwise, mixed-precision bfloat16 training was impossible due to OOM. In this setting using two 80GB GPUs, pure bfloat16 was 39.2% faster than using mixed precision. In our setting with four parallel 80GB GPUs, pure bfloat16 was 31.0% faster than using mixed precision. We note that a large part of this speedup is mainly enabled by not having to apply activation checkpointing to fit the model during training and delaying the syncing of gradients during gradient accumulation. When comparing pure bfloat16 using the same training configuration as the most efficient mixed-precision run, pure bfloat16 was only 11.4% faster.8 We also benchmark using eight parallel GPUs for reference, although this is arguably stretching a tight academic compute budget. Here, pure bfloat16 is also 29.8% faster than mixed-precision training. Although we do not have access to sufficient compute resources to benchmark this in our setup, with sufficiently many parallel devices, mixedprecision training does not require activation checkpointing to efficiently train large-scale models, eliminating one of the main benefits of pure bfloat16. This highlights the benefits of pure bfloat16 especially in settings with limited total GPU memory due to a lower number of parallel devices. In such settings, mixed-precision training is either impossible or requires computationally expensive memory-saving techniques, which pure bfloat16 does not require to the same extent.NLL Loss and downstream task results. Comparing the loss over the course of continued pretraining of pure bfloat16 and mixed-precision bfloat16 training in Table 3 , we see that pure and mixed-precision bfloat16 training are very close. However, at the very beginning (after 10% of training steps) and towards the end (from 90% of training steps), we find that mixed precision achieves a slightly lower loss. This is very likely an artifact of our cosine training schedule with linear warmup: at the very beginning and end, the learning rate is very small and results in small weight updates, which are problematic in pure bfloat16 training as we will discuss below. We also compare the downstream task performance of pure bfloat16 and mixed-precision bfloat16 training in Table 5 . Interestingly, pure bfloat16 outperforms mixed-precision bfloat16 in our downstream task evaluations both in German and English. Note that pure bfloat16 training effectively acts as a regularizer that zeroes out updates in the optimizer step that are smaller than some ϵ dependent on the magnitude of the to-be-updated weight according to the formula outlined in Section 2. We analyze this below.The effects of pure bfloat16 numerics. As shown in Figure 2 , the parameter values of RMSNorm layers in Mistral-7B are particularly large, which results in a larger ϵ below which updates are squashed to zero in pure bfloat16 training. We further investigate this in bfloat16 numerics, the RMSNorm layers receive almost no updates in pure bfloat16 training due to their larger weight values. We see that pure bfloat16 also reduces the average parameter change for other layers, but this effect is much less pronounced.It is not clear that the regularizing effect of pure bfloat16 training is a desired behavior, even though it results in better downstream task performance in our experiments. However, given the large training efficiency boost possible through pure bfloat16 training when on a tight academic compute budget, we believe that this potentially undesired sideeffect is a worthwhile tradeoff to consider. We stress that we perform a data-matched comparison between pure and mixed-precision bfloat16. A compute-matched comparison, where pure bfloat16 will train for significantly more steps, would be even more favorable for pure bfloat16 training. Additionally, we can enable mixed precision or float32 just for the parameters with large values (in the case of Mistral-7B these are the RMSNorms) to offset the majority of bfloat16 regularization impact.In terms of loss, we do observe slightly better performance of mixed precision at the very end and beginning of training due to the linear warmup and cosine decay of the learning rate. As the final annealing phase has been shown to be important for model performance (Ibrahim et al., 2024) , we further study this in our hindsight runs, which are trained in pure bfloat16. We employ an "infinite" learning rate schedule (Ibrahim et al., 2024; Zhai et al., 2022) with a prolonged constant learning rate before the final annealing down to a smaller value. We use the same checkpoint before the beginning of the annealing phase and continue training with either pure or mixed-precision bfloat16. Using mixed precision during the annealing phase is denoted as pure++ in Table 5 and Table 7 . Interestingly, we do not see any conclusive advantage of using mixed precision during the annealing phase. We conclude that fully pure bfloat16 continued training is indeed viable even during the low learning rate annealing phases of a learning rate schedule.The next research question we investigate in this work are the benefits of tokenizer swapping for language adaption of pretrained LLMs. This question deserves special attention in the tight academic budget setting, as there is less compute spent on re-learning the new embeddings.On comparing loss between different tokenizers. Comparing the loss (negative log-likelihood) across different tokenizers is not straightforward. Different tokenizers produce different numbers of tokens for the same text. The conventionally reported loss is the summed negative loglikelihood normalized by the number of tokens -a direct advantage for worse tokenizers that produce more tokens for the same text.Instead, we choose to normalize by a tokenizer-independent quantity: the number of (white-space split) words. This choice is arbitrary in some respects as it is just a constant for all methods. Instead, we could also normalize by the number of characters or bytes (see Gao, 2021; Forsythe, 2023; Gao et al., 2020) . Choosing words as the normalizing factor has the advantage of yielding a similar value range to the familiar token normalization. The word-normalized negative log-likelihood is computed over the same chunks of text for each model. We report the token-normalized log-likelihood in Appendix C.Embedding re-initialization. When swapping tokenizers for language adaptation, it has been shown that a good re-initialization of the new embedding matrix is crucial (Minixhofer et al., 2022; Ostendorff & Rehm, 2023; Dobler & de Melo, 2023; Downey et al., 2023) . We evaluate various embedding initialization methods and find that FO-CUS (Dobler & de Melo, 2023) performs best for German.For Arabic in our hindsight study, all methods showed suboptimal performance, so we use FOCUS coupled with an additional short gradient descent training phase of the embeddings with the rest of the model frozen for 100 steps (roughly 1.5% of total training steps). Full details are provided in Appendix B.Results of tokenizer swapping. In Table 3 We report downstream task results on a suite of German benchmarks in Table 5 . We do not see a clear trend of better performance with or without tokenizer swapping. Note that for inference on downstream tasks in the target language, however, tokenizer swapping will be computationally more efficient, as the new tokenizer will produce fewer tokens for the same text (Yamaguchi et al., 2024) . 9 In Table 6 , we also report the average performance on English downstream tasks. As expected, we find that replacing the original tokenizer with an exclusively German tokenizer leads to diminished results on English. Thus, whenever performance in the source language is important, the new tokenizer should perhaps retain a larger share of the original tokens.The base model's tokenizer already contains some German tokens10 , as well as many English tokens also used in German. Our experiments hence provide a rough lower bound on the effectiveness of tokenizer swapping. We expect much larger benefits for unseen and low-resource languages. Even though downstream tasks on German do not improve significantly with a specialized tokenizer, our results demonstrate the viability of re-learning an embedding matrix after tokenizer swapping even when on a tight academic compute budget. In contrast, Zhao et al. (2024b) Table 7 : Results on Arabic downstream task suites. The average is a macro-average that includes the individual benchmarks in AlGhafa-T. † : For pure++ bfloat16, mixed precision was used just for the final annealing phase of the learning rate schedule.derperforms compared to keeping the original vocabulary for tight compute budget continued pretraining on Chinese. However, we note that our experiments for tokenizer swapping are compute-matched but not data-matched: since the new tokenizer encodes the training data more efficiently, more total text was seen during continued pretraining for the same number of tokens. Our experiments do not suggest that this increased efficiency translates to better performance.From the results in Table 6 , we have to conclude that the resulting models from our main experiments do not show good downstream task performance. Strikingly, the base Mistral-7B model achieves better downstream task results than our adapted German variants, even on German benchmarks. Mistral-7B has already seen German during pretraining and its tokenizer contains German tokens, which allows us to roughly lower bound language adaptation effectiveness, but for most other languages it would be a pessimistic estimate. Therefore, we further adapt Mistral-7B to Arabic alongside German in an additional hindsight study. Also, the models trained in our hindsight study were trained using higher quality data from CulturaX rather than OSCAR23.01 and used several other improved training techniques (see Section 3.2).The German hindsight study model using the custom German tokenizer has an average improved downstream task performance of two percentage points over the best model from the main experiments using the exact same computational budget. This demonstrates the importance of data quality and our training improvements such as preventing cross-document attention contamination via an adjusted attention mask. However, the model still performs worse than the base Mistral-7B and its adapted German LeoLM version on German downstream tasks.Arabic downstream results. In contrast, the Arabic hindsight study models perform significantly better than the base Mistral-7B model on Arabic benchmarks (reported in Table 7 ). Interestingly, even though the goal of our study is analysis rather than achieving state-of-the-art, our Arabic models also outperforms AceGPT-7B (Huang et al., 2023) , which is a continued pretraining of Llama 2 on Arabic for 3.75× more tokens than our models but without tokenizer swapping. Note that the training recipes for our German and Arabic hindsight study models were exactly the same -highlighting that language adaptation is especially useful for underrepresented languages.When is language adaptation on a tight academic compute budget viable? Although the training data of Mistral-7B is not open, due to its multilingual performance and the existence of German tokens in its tokenizer, we can conclude that it was pretrained multilingually and that German had a significant share. Intuitively, specializing such a model on a single language might boost performance, e.g., by removing curse of multilinguality (Conneau et al., 2020) capacity bottlenecks. Our experiments with German as a target language show that continued pretraining to focus model capacity just on the target language is not always beneficial. This is likely because German was already well-represented.Recent language adaptation attempts for lesser-resourced or less well-represented languages than German, e.g., Polish (Ruciński, 2024) or Chinese (Zhao et al., 2024b) , do report improvements in the target language compared to the base model with similar computational budgets. Le-oLM used eight times as much compute as our study and obtains a minor increase in German task performance of 0.6 percentage points.Our results on Arabic highlight that target languages that were not dominant or unseen during pretraining can indeed benefit from language adaptation on a tight academic compute budget. In our setup, we focus only on monolingual performance in the target language and do not include English or code in the training data. Recent work has shown that the inclusion of such data can boost performance even in the target language (Muennighoff et al., 2023; Csaki et al., 2023) . However, it is not clear if such benefits will also materialize on a tight academic compute budget. 
related work: 
 Language adaptation of pretrained LLMs. Adapting strong pretrained LLMs such as Llama (Touvron et al., 2023a; b) or Mistral-7B (Jiang et al., 2023b) to other languages has recently had a surge in popularity (e.g., Plüster et al., 2023; Pires et al., 2023; Cui et al., 2023; Nguyen et al., 2023; Alves et al., 2024) . However, these studies use significantly more compute than our setting and many either do not specialize the tokenizer or do not use strong embedding initialization methods. Csaki et al. (2024) adapt pretrained LLMs to a variety of languages and find that specializing tokenizers through vocabulary expansion improves fertility but not downstream task results. Recent work has shown that mixing in English during continued pretraining helps downstream task performance in the target languages (Csaki et al., 2023; Zhao et al., 2024b) , whereas Ibrahim et al. (2024) find that skewing the data mix towards the target language helps adaptation at the cost of more forgetting. Muennighoff et al. (2023) show in data-constrained experiments on English that individual samples can be repeated up to four times and mixing in code data can improve downstream task performance. 2020) train Transformer models from scratch in pure bfloat16 and find that stochastic rounding of the weight updates was necessary to recover the performance of using float32 optimizer states and master copy of the weights. Despite this, the most popular choice for training language models in many popular libraries remains mixed-precision bfloat16 (Lewandowski & Kosson, 2023) . We show that for continued pretraining, pure bfloat16 training is viable without any further modifications such as stochastic rounding. 
conclusion: 
 Based on our results, we draw the following final conclusions.Pure bfloat16. Training in pure bfloat16 (rather than using mixed precision) is viable for continued pretraining and enables significantly faster training when only a few parallel devices are available. Pure bfloat16 training does come with certain caveats: We find that for Mistral-7B, the RMSNorm weights (which are larger than most other weights) are not updated in pure bfloat16 training due to the numerics of bfloat16. However, this has no clear drawbacks in terms of loss and downstream task performance.Tokenizer swapping. Tokenizer swapping coupled with a good embedding initialization is viable and performs at least on par with keeping the original vocabulary but did not significantly improve downstream task results for German. We note that Mistral-7B's tokenizer already contains some German tokens. For Arabic, which was an unseen script in the original tokenizer, our adapted model with tokenizer swapping fares significantly better than the base model without language adaptation. Our results demonstrate that re-learning the embeddings for a new tokenizer is viable even when on a tight compute budget.Tight compute budgets for language adaptation. Adapting Mistral-7B to German on a tight compute budget performed worse than the base Mistral-7B for German downstream tasks, irrespective of tokenizer swapping and training precision. Adapting to Arabic however significantly increased Arabic downstream task performance. This highlights that language adaptation is not always beneficial but can be, especially if the target language was not wellrepresented. We believe that our findings regarding pure bfloat16 training and tokenizer swapping can help inform future language adaptation efforts towards more efficiency. 
 
 Appendix: We limit our study to a single LLM (Mistral-7B). This is motivated by both the checkpoint's popularity and the popularity of its parameter-size class (seven billion parameters). However, it is not certain that model-specific findings, such as RMSNorm updates being flushed to zero in pure bfloat16 training, will transfer to other models. Our general findings on weight updates in pure bfloat16 to large values are expected to generalize. Also, while we add Arabic as a second language in our hindsight study, we consider only German and Arabic in our experiments. Although less significant for our findings on pure bfloat16, the efficacy of tokenizer swapping will likely vary depending on the target language.Additionally, a large part of the training efficiency gains of pure bfloat16 training compared to mixed-precision training materialize mainly when the training is constrained by GPU memory. For models smaller than Mistral-7B that might not require sharded training, these benefits will be less pronounced.We note that we analyze the viability of pure bfloat16 only for continued pretraining, where the weights effectively start with an initialization (reasonably) close to their optimized value. We cannot draw the same conclusion for pretraining from scratch, as the training dynamics in the initial phase of pretraining from randomly initialized weights might have a much larger dependence on using high-precision optimizer states and a master copy of the model weights to reduce numerical errors.We consider only full-parameter training and do not evaluate the use of techniques such as LoRA (Hu et al., 2021) We publish our training code at https://github.com/ konstantinjdobler/tight-budget-llm-adaptation.We use PyTorch FSDP for sharded training, and we rely on the lightning package for the implementation of distributed and mixed-precision training. We employ document packing and concatenate samples separated by an [EOS] token until the block size is filled. In our hindsight runs, we instead prepend a [BOS] token and adjust the attention mask so that cross-document attention is prevented. We enable TF32 computation for float32 via torch.set_float32_matmul_precision("high").For pure bfloat16 training, the Mistral-7B implementation from HuggingFace implements some small parts of the forward pass in float32, which we do not modify. Specifically, these are the input variance calculations in the RMSNorm layers and the final softmax in the language modeling head. All weights and optimizer states are still in bfloat16 for pure bfloat16.For the "infinite" learning rate schedules in our hindsight study, we follow Ibrahim et al. (2024) and use a 1% warmup up to a maximum learning rate of 3 × 10 -5 , then a cosine decay for 60% of steps to 1.65 × 10 -5 , followed by a constant learning rate for 25% of steps, and finally annealing down to 2 × 10 -6 for the last 14% of steps.We use Docker images to conduct our training in a fullyreproducible environment, which are published alongside lockfiles with the pinned package versions that we used.When swapping tokenizers for language adaptation, it has been shown that harnessing the compute spent on the original embedding matrix via a good re-initialization of the new embedding matrix is crucial (Minixhofer et al., 2022; Ostendorff & Rehm, 2023; Dobler & de Melo, 2023; Downey et al., 2023) . We evaluate FOCUS (Dobler & de Melo, 2023) and WECHSEL (Minixhofer et al., 2022) , two initialization methods based on mappings to the old embedding space. Furthermore, we evaluate a heuristics-based approach (Downey et al., 2023) , initializing from N (original_embs_mean, original_embs_std) or N (0, 0.02), and simply copying the old embeddings, effectively randomly assigning old embeddings to new tokens. 11We evaluate all the mentioned methods directly after initialization without any further training and report the results in Table 8 . We find that out of these methods, FOCUS provided the best results directly after initialization. Therefore, we utilize FOCUS for all experiments with tokenizer swapping. Note that these results might differ for different target languages and the gap between these methods is reduced by continued pretraining.For Arabic in our hindsight study, all methods showed suboptimal performance, so we use FOCUS coupled with an additional short gradient descent training phase of the embeddings with the rest of the model frozen for 100 steps (roughly 1.5% of total training steps). Specifically, FO-CUS likely underperformed because we deliberately did not add English tokens to the new tokenizer, leading to a very small -and "low-quality" overlap with the original vocabulary, which FOCUS relies on. This could be ameliorated by including more tokens from the original vocabulary or employing a bilingual dictionary.Very recently -after our experiments had finished - Minixhofer et al. (2024) proposed ZeTT, which utilizes a trained hypernetwork to predict embeddings for a new tokenizer. We briefly discuss this method since it yielded very promising initial results. We found that ZeTT achieves a lower initial loss than FOCUS for German and Arabic. For Arabic specifically, it is the only method that performed significantly better than random. ZeTT requires a hypernetwork trained specifically for the base model, which is not computationally sensible as a one-time effort for the embedding initialization. However, the authors do provide such a hypernetwork for Mistral-7B and select other popular base models. The computational cost of predicting the new embeddings using the hypernetwork is negligible. If such a hypernetwork is available, using ZeTT could further improve the performance of tokenizer swapping over our current results using FOCUS.Alongside the word-normalized negative log-likelihood reported in Table 3 in the main body of the paper, we report the token-normalized negative log-likelihood (the conventional cross-entropy loss) in Figure 3 . Note that this is not a We compare pure and mixed-precision bfloat16 training. Additionally, we compare swapping the original tokenizer of Mistral-7B with a specialized German tokenizer.fair comparison and we only provide this for completeness, as Mistral-7B's original tokenizer will produce more tokens for the same text, which directly results in a lower loss as the loss is divided by the number of tokens.We report the full downstream task results for all models on German (in Table 9 ) and on English (in Table 10 ).For German benchmark evaluation, we use the benchmark suite provided by Plüster (2023) , which was also used for LeoLM. We use their fork of the lm-eval-harness (Gao et al., 2023) to evaluate our checkpoints. The individual datasets used are translated versions of MMLU (Hendrycks et al., 2021) , TruthfulQA (Lin et al., 2021) , ARC (Clark et al., 2018a) , HellaSwag (Zellers et al., 2019) , as well as PAWS-X (Yang et al., 2019) and LAMBADA-OpenAI (Paperno et al., 2016; Radford et al., 2019) .For the Arabic benchmarks based on OALL (Elfilali et al., 2024) , we used ACVA (Huang et al., 2023) , AlGhafa (Almazrouei et al., 2023) , and the following translated benchmarks from AlGhafa-T: MMLU (Koto et al., 2024; Hendrycks et al., 2021) , EXAMS (Hardalov et al., 2020) , ARC-Challenge (Clark et al., 2018b) , ARC-Easy (Clark et al., 2018b) , BOOLQ (Clark et al., 2019) , COPA (Roemmele et al., 2011) , HellaSwag (Zellers et al., 2019) , Open-BookQA (Mihaylov et al., 2018) , PIQA (Bisk et al., 2020) , RACE (Lai et al., 2017) , SciTail (Welbl et al., 2017) , TOX-IGEN (Hartvigsen et al., 2022) . We report the full results of Arabic models for all benchmarks in Table 11 . The Arabic benchmarks were evaluated using lighteval (Fourrier et al., 2023) in a 5-shot setting.We ran a performance benchmark inspired by Hagemann et al. (2023) to determine the best possible combination of micro-batch size, activation checkpointing, FSDP ZeRO sharding, syncing during gradient accumulation (not syncing is faster but incurs a memory overhead), and utilizing a paged variant of AdamW for each combination of number of GPUs and precision type. We evaluate the full Cartesian product of micro-batch sizes {1,2,4,8}, activation checkpointing {yes, no}, FSDP ZeRO sharding {full, grad_op}, gradient accumulation syncing {sync, no_sync}, and AdamW implementation {paged, no_paged}. For training on a single GPU, mixed-precision training is impossible even with all memory saving techniques maximized. For pure bfloat16 on a single GPU, we instead manually search for the best configuration. We do not consider FSDP CPU offloading, as this did not yield acceptable compute efficiency. We report the full results of this benchmark in Table 12 . For the purposes of this paper, we actually run the full benchmark to compare pure and mixed-precision bfloat16 using the same configuration, even if a better configuration is also available for pure bfloat16. We do make one small optimization and exclude trials where a more memory-friendly configuration previously went OOM. To save compute, one could additionally prune benchmark trials through educated guesses of suboptimal configurations (e.g., when using eight devices, we do not need to benchmark heavy memory optimizations) or employ a binary instead of linear search. For all benchmarked configurations, we run 11 steps and exclude the first one from the reported average since it is disproportionately affected by initial one-time costs. Due to the way we extract step timings from our original training script, timings are rounded to the nearest second for step times over a minute instead of considering two decimal places.MMLU HellaSwag TruthfulQA ARC LAMBADA PAWS-X Avg.German Hasso Plattner Institute / University of Potsdam, ELLIS Unit Potsdam, Germany. Correspondence to: Konstantin Dobler <konstantin.dobler@hpi.de>. Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML2024).Typically, unless one is willing to exchange significant social goodwill into compute by delaying the work of all other people in the lab, multiple GPUs can be blocked for only a short number of days (e.g., over weekends or holidays).In this work, Mistral-7B specifically refers to mistralai/Mistral-7B-v0.1.Recently, the fp8 type supported by Nvidia H100 GPUs has also gained popularity, but we leave this for future work. The problems we illustrate for bfloat16 are expected to be exacerbated for fp8.We observe very large gradient norm spikes early on when training with 6 × 10 -5 , 2 × 10 -4 , and 3 × 10 -4 .We remove samples with any of the following quality warnings: adult, noisy, header, footer, tiny, short_sentences. The quality warnings are provided by the OSCAR corpus.By this, we refer to all states necessary during training: model weights, gradients, optimizer states, and activationsUsing torch_model.no_sync() during gradient accumulation is a common optimization to reduce communication overhead but results in larger maximum memory usage when coupled with FSDP since the gradients will not be sharded for that period.8 See the full benchmark results in Appendix E.This can lead to significant computational savings: In our evaluation on Arabic, Mistral-7B took over 5 times as long as a model with our custom Arabic tokenizer due to poorer tokenizer fertility.Such as _Jahrhunderts or meisterschaft.Since our new tokenizer has a larger vocabulary than the original one, we initialize the remaining tokens from N (0, 0.02). 


Decision: We thank the authors for their time and contribution to WANT and we are pleased to share that after the reviewing process the paper has been accepted. Congratulations! We encourage the authors to consider reviewers' feedback for the improvement of the camera-ready version. We hope to see you in person at the workshop and brainstorm on efficient training research together!