nanoLM: an Affordable LLM Pre-training Benchmark via Accurate Loss Prediction across Scales

LLMs’ training involves cross-entropy objectives such as predicting subsequent or masked tokens. Since typically the training data is processed for less than 1 epoch, this loss also strongly correlates with the perplexity on validation set. Besides, recent researches (Touvron et al., 2023b; Schaeffer et al., 2023b) suggest a link between pre-training loss and downstream task performance, although this is challenging to assess comprehensively. In short, for the same data, pre-training loss is a reasonable indicator of model capabilities.

We explore the problem of predicting the loss of a target (large) model with a series of small proxy models. Based on the discussions in Sections 2.2 and 2.3, we observe that in this problem setting, $\mu$P and scaling laws are complementary: first, scaling law offers a formal function to directly calculate the loss values with model scales, but the optimal HP of each model should be searched separately; second, $\mu$P provides a unified HP for multiple models varying only in width, but does not guarantee the exact loss values except knowing “the wider, the better”. Thus, we develop a loss prediction approach, namely $\mu$Scaling, as follows (Figure 2, bottom):

First, we generate a series of proxy models $\{M(w_{0}),M(w_{1}),M(w_{2}),...,M(w_{n})\}$ with $M_{B}\!=\!M(w_{0})$ being the base model and $M_{T}\!=\!M(w_{n})$ being the target large model. They differ only in widths $\{w_{0},w_{1},...,w_{n}\}$.

Second, we grid-search for the optimal $\mu$Transferable HPs $H$ = (learning rate, init variance, multipliers) on $M_{B}$. $\mu$P concludes that $H(i)\!=\!\mu P(H(0),w_{i}/w_{0})$ are also the optimal HPs for each $M(w_{i}),i\!\in\![0,n]$. We illustrate the $\mu$P function we use in Table 2, which corresponds to Table 8 in (Yang et al., 2022).

Third, we select some small $w_{i}$-s, train the models $M(w_{i})$ with $H(i)$, and achieve loss $L(i)$. Based on these points, we fit a power-law function $L(C)=aC^{b}+c$ w.r.t. the number of parameters $C$, and expect that this function accurately predicts $L_{T}\!=\!L(n)$ that could be achieved by training $M_{T}\!=\!M(w_{n})$ with HP $H_{T}\!=\!H(n)$.

The correctness of the third step is critical and has not been studied yet, but we show by our comprehensive experiments (Section 4) that it is effective for HPs within the common loss basins in the HP space, and with our modification on scaling laws (below). Note that the formulation of scaling law is empirical and should be validated by experiments (Kaplan et al., 2020; Hoffmann et al., 2022; OpenAI, 2023).

Embedding counts as model size. An interesting difference between $\mu$Scaling and mainstream scaling laws is we find that $\mu$Scaling fits better if the size of embedding is counted in the model size, while methods like (Kaplan et al., 2020) benefit from excluding it. We explain this in Section 4.4.

Based on $\mu$Scaling, the $\mu$Transferable HPs are actually “tunnels” towards directly comparing different model designs on large scales via loss prediction. For each model, we predict its loss via $\mu$Scaling by shrinking its width into a series of small ones and do not need to actually train the large model with full width. In this process, only one HP search is needed for each model design. We formalize this new paradigm in Algorithm 1, and highlight its feature as “research without re-search”. This makes LLM studies more affordable, while still reaching meaningful conclusions for model scales beyond 50B.

Our main goal is to justify the third step in $\mu$Scaling (Section 3.1), as well as the usability of the nanoLM benchmark. To validate the loss prediction capability with different data scales, models, and computational resources, we experiment with the following three separate settings:

Single GPU. Firstly, we test loss prediction on Transformer architectures (i.e., GPT, BERT, T5) with 12 layers using the C4 (Raffel et al., 2020) dataset under a single-GPU setup. We utilize a base width of 256 to conduct a grid search for the $\mu$Transferable HPs, primarily targeting the learning rate, with values ranging from $5e-4$ to $1e-1$^†††The $\mu$Transferable HPs include learning rate, initialization standard deviation, and multipliers. Due to computational resources, we focus on the learning rate. However, researchers can still explore the other $\mu$Transferable HPs on nanoLM to achieve even better results.. The number of parameters for the model series ranges from 8M to 700M with different widths, as detailed in Table 3(a). The target width for prediction is 1024 ^‡‡‡Given that this group of experiments target on only one single GPU, both BERT and T5 ran out of memory when expanded to a width of 2048. In these cases, we leave the results for future.. The batch size was established at 16, and the sequence length was set to 2048. According to (Yang et al., 2022), when the training steps surpass 5k, the $\mu$P transfer tend to stabilize. To this end, we choose to fit the loss at 7k and 20k steps and study its influences on the predicative performances.

Multi-GPU with FSDP. Secondly, to experiment with larger models and more data, We employ Fully Sharded Data Parallel (FSDP) (Zhao et al., 2023) to effectively address the training challenges. We conduct loss prediction with 32-layer GPTs on nanoLM pre-training data (Section 3.2). We utilize a base width of 256 to conduct a grid search for the learning rate, ranging from $5e-4$ to $1e-1$. The number of parameters ranges from 38M to 26B (Table 3(b)). The batch size was established at 512, and the sequence length was set to 512. We use proxy models with widths under 2048 to predict the loss of the target model with width 8192.

Multiple Machines with Megatron. Lastly, we further validate the feasibility of $\mu$P and $\mu$Scaling with extremely large widths. We conduct the experiments for the 64-layer GPTs on Megatron (Narayanan et al., 2021). It provides efficient tensor, pipeline, and sequence-based model parallelism for pre-training LLMs. We utilize a base width of 256 to conduct a grid search for the learning rate, ranging from $1e-4$ to $1e-2$. The number of parameters ranges from 77M to 52B (Table 3(b)). The batch size was established at 512, and the sequence length was set to 2048. We exit training and fit the loss at 10k steps, with 10.49B tokens consumed (10$\%$ of the 100B track of nanoLM pre-training data, Section 3.2).

All of our experiments are run on A100 GPUs. For each series of models used for loss prediction, the batches are fed into the models in the same order. More hyperparameters can be found in Appendix C.

All of our models are trained from scratch. Additionally, we adhere to the recommendation by (Yang et al., 2022) to initialize the output word embeddings and query parameters as all-zero to avoid the Gaussian Process ^§§§Such Gaussian Process may cause misalignment of landscapes between small and large models.. We maintain a consistent head dimension of 64 for each attention head and scale up the number of heads with model widths. The dropout and weight decay are set to 0. We employ the AdamW optimizer (Loshchilov & Hutter, 2019) with its default configurations. The ${a,b,c}$ coefficients in power law $L=aC^{b}+c$, as well as their standard deviations are computed with $scipy.optimize.curve\_fit()$ ^¶¶¶https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html.

For GPT, Llama, BERT, and T5, proxy models with widths $\leq$ 896 ^∥∥∥We found a positive correlation between the number of data points and the accuracy of the fitted loss curve. We choose seven points due to the limited computation resource. are used to fit the power-law curves and predict the loss for widths $>$ 896. For the 12-layer GPT and Llama, the fitted power-laws are $L=2.07*C^{-0.52}+3.85$ and $L=1.34*C^{-0.49}+4.02$, respectively. For 12-layer BERT, the result is $L=61.18*C^{-1.31}+3.43$; for the 12-layer T5, it is $L=0.25*C^{-0.47}+2.82$. According to Figure 3, it’s evident that based on the loss of the narrow proxy model series, we successfully predicted the loss of wider models. This holds for both 7k and 20k training steps and for all the 4 different Transformer variants, supporting the effectiveness of $\mu$Scaling. This extends the $\mu$P conclusion (Yang et al., 2022) of “the wider, the better” to “wider is better adherent to a power law”. In contrast, for the yellow points, whose HPs are not transferred by $\mu$P but follows standard parametrization, the scaling law is not well fitted. They also underperform the same-width networks using $\mu$P. Additionally, Figure 3 (d) demonstrates that the conclusion still holds on diverse multilingual datasets, such as MC4 (Raffel et al., 2020).

To validate $\mu$Scaling for larger-scaled models, we expanded the GPT to 32 layers. We train with the data parallel strategy (FSDP) on the benchmark pre-training data, as supported by nanoLM (Section 3.2). Based on the proxy models with widths ranging from 256 to 2048, we predict the loss for a model with a width of 8192. The results are presented in Figure 4 with the power-law being $L=0.077*C^{-0.61}+3.37$. Notably, $\mu$Scaling successfully predicts the loss of a 26B model at 7k steps based on the losses from models ranging from 38M to 1.7B. This demonstrates that nanoLM is an effective paradigm to study LLMs of billion-scale magnitude.

We further validate nanoLM with extremely large models, exemplified by 64-layer GPTs implemented with Megatron (Shoeybi et al., 2019). We train proxy models with widths ranging from 256 to 2048 using nanoLM pre-training data and forecast the loss for a model in 8192 width. Results are presented in Figure 4 with the power-law being $L=0.249*C^{-0.47}+2.82$. The results show that we can successfully predict the loss of a 52B model at 10k steps based on proxy models ranging from 77M to 3.4B, further affirming the accuracy of $\mu$Scaling and usability of nanoLM.

We calculate the number of floating point operations (FLOPs) of a model following (Narayanan et al., 2021). Given a Transformer model with $l$ layers, width (a.k.a. hidden size) $w$, sequence length $s$, vocabulary size $V$, and training batch size $B$, the total number of floating-point operations is $M(l,w,s,V,B)=96Bslw^{2}\left(1+\frac{s}{6w}+\frac{V}{16lw}\right)$. Thus, the FLOPs ratio of the whole nanoLM process to the training of the full target model can be approximated as:

where $w_{1}$ is grid-search width and $t$ denotes corresponding number of grid search trails on $\mu$Transferable HPs (as detailed in Table 2). $n$ signifies the count of distinct proxy models used for the fitting loss function, while $w_{t}$ is the target model width. In our experiments, the 32-layer GPT comes with a sequence length of 512 and a batch size of 512. The vocabulary contains 100,256 tokens. Thus, the ratio of cost is approximated at 0.131. Under similar calculation, the cost ratio for our 64-layer GPT is 0.142. In short, with just $13.1\%,14.2\%$ of the one-time pre-training cost, we can predict the loss for 26B and 52B models. Note that without $\mu$Scaling, grid-search is needed on the large models, which will enlarge the advantage of our method by several times.

We compare the fitting error of nanoLM with related work that also provided loss numbers, including Cerebras-GPT (Dey et al., 2023) (in their Table 8), and Pythia (Biderman et al., 2023). As shown in Table 4, for fitting error, nanoLM achieves slightly better results than related work, even with larger models (26B and 52B). Moreover, the covariances of the fitted coefficients $\{a,b,c\}$ are significantly lower than related work. This indicates that nanoLM is more accurate and confident on its scaling laws, supporting its reliability.

We demonstrate an ablation study by Figure 4(a) and 4(b), indicating that our scaling law fits worse if embedding sizes are not counted as model sizes, in contrast to (Kaplan et al., 2020; Henighan et al., 2020). This is potentially because $\mu$P concludes that the learning rate of embedding layers should not be scaled down with widths, while most existing methods grid-search for a unified learning rate for all layers on each width, making embeddings learning slower on large models, and the matrix-like parameters dominating the training dynamics. This analysis is conducted with a 12-layer GPT-2 trained on OpenWebText (Gokaslan & Cohen, 2019) for 20k steps.

For other discussions, please see Appendix A.