4 Pretraining

本文为论文《SmolLM2:紧凑高效的小语言模型家族》「论文原文」板块的对应章节,按 arXiv 官方全文逐字保留,层级与「中文全文翻译」板块一一对应。

Recent trends in language models pretraining show a clear shift towards significantly longer training durations, especially for smaller models (Yang et al., 2024a; Yang et al., 2024b; AI@Meta, 2024b). While this strategy deviates from the Chinchilla-optimal guidelines (Hoffmann et al., 2022), the resulting performance gains and reduced inference costs make extended training a worthwhile trade-off (de Vries, 2023). For example, Qwen2-1.5B was trained on 7 trillion tokens, Qwen2.5-1.5B on 18 trillion tokens, and Llama3.2-1B, derived from a pruned 8B model, was trained using distillation on 9 trillion tokens (Yang et al., 2024a; Yang et al., 2024b; AI@Meta, 2024b).

When building SmolLM2, we trained on 11 trillion tokens (approximately two epochs on our collected datasets), employing a multi-stage training approach instead of a fixed dataset mixture throughout pretraining. This design was guided by four key principles: (1) Performance-driven interventions, where we monitor evaluation metrics on key benchmarks and adapt dataset mixtures to address specific capability bottlenecks; (2) Upsampling high-quality math and code during the annealing phase, reserving datasets like FineMath and parts of Stack-Edu for the final stages to maximize their impact (Blakeney et al., 2024; Ai2, 2024); (3) Strategic introduction of medium-sized datasets, such as OWM, InfiMM-WebMath, and Stack-Edu, mid-training to avoid dilution by larger datasets early on; and (4) Avoiding excessive data repetition, in line with Muennighoff et al. (2023) we aimed to stay close to the recommended 4–5 epoch threshold for most datasets. While it might be fruitful to perform multiple from-scratch training runs to explore different data mixing schedules, the high cost of pretraining SmolLM2 (around $250,000 USD of GPU compute) motivated our “online” approach.

In the following sections, we describe each stage of the training process, detailing the dataset mixtures, the rationale behind our choices, and the observations that guided our interventions. While some decisions were informed by established findings in the literature, others were driven by empirical insights gathered during training. The data mixtures of the four pretraining phases are available in Figure 2.

4.1 Training setup

Our base model contains 1.7B parameters and follows the LLama2 (Touvron et al., 2023) architecture, outlined in Appendix A. We trained the model on 256 H100s using the nanotron framework and use AdamW optimizer with $(\beta,\beta_{2})=(0.9,0.95)$ with a Warmup Stable Decay (WSD) (Hu et al., 2024; Zhai et al., 2022) learning rate schedule to avoid setting a fixed training duration (see Figure 3, Appendix A). The schedule started with a 2,000-step warmup phase, maintained a peak learning rate of $5.0\times 10^{-4}$ (stable phase), and could transition to a decay phase when needed, reducing the learning rate to zero over 10% of the total training steps (Hägele et al., 2024). We used the tokenizer from Allal et al. (2024), which has a vocabulary size of 49,152 tokens and was trained on a mixture of 70% of FineWeb-edu, 15% Cosmopedia-v2, 8% OpenWebMath, 5% StarCoderData and 2% StackOverflow.

图

Figure 2: Dataset mixtures across training stages. Detailed descriptions are provided in Section 4. The x-axis represents the number of training tokens.

4.2 Stable phase: stage 1

Data mixture

In the first phase of SmolLM2’s pretraining (0 to 6T tokens), we designed our dataset mixture based on insights from our English web ablations and existing literature. We adopted a 60% FineWeb-Edu and 40% DCLM ratio (discussed in Section 2.2) for web data, which provided an optimal balance between educational content and diverse, real-world Q&A-style data. For code data, following Aryabumi et al. (2024), we incorporated StarCoderData, consisting of 250B tokens across 80 programming languages, and limited it to 10% of the total mixture to ensure approximately 4 epochs over 11T tokens with room for upsampling in later stages. We did not include math data in stage 1 due to our math datasets’ relatively small size.

Findings

After 6T tokens of training, we evaluated SmolLM2 on key benchmarks, as shown in Table 3. Knowledge and reasoning performance aligned with expectations based on our English web ablation results. However, we observed generally poor coding and mathematics performance.

4.3 Stable phase: stage 2

Data mixture

For stage 2 (6T to 8T tokens), we added OWM to the mixture at a 5% ratio and increased the proportion of code data in hopes of maintaining strong knowledge retention while addressing observed gaps in coding and mathematical reasoning. Including OWM at a low percentage reflects the dataset’s small size (12B tokens) and our gradual approach to incorporating math content. The final mixture for stage 2 consisted of 75% English web data (keeping the 60/40 FineWeb-Edu to DCLM ratio from stage 1), 20% code data, and 5% math data, as shown in Figure 2.

Findings

After stage 2, code performance improved across most languages, validating the decision to upsample StarCoderData. OWM integration had no significant impact on math performance, underscoring the need for larger, higher-quality math datasets in later stages. Beyond code and math performance, as shown in Figure 6 (Section E.2), we observed above-random (>25%) MMLU accuracy with a multiple-choice formulation (MCF, i.e. explicitly outputting an option from ’A’, ’B’, ’C’, or ’D’ instead of computing the likelihood of different answers as in the cloze formulation). This contrasts prior work showing that small models struggle with the MCF (Gu et al., 2024; Du et al., 2024) and suggests that long trainings of small models can make them acquire abilities typically associated with larger models (Blakeney et al., 2024; Gu et al., 2024; Du et al., 2024). To further optimize MMLU performance, we revisited our English dataset mixture with additional annealing ablations and found that increasing DCLM relative to FineWeb-Edu slightly improves MMLU MCF at this stage.

Table 3: Average model performance on different benchmark categories after each training stage. Stages 1-3 are during stable phase (no decay). Full per-benchmark results in Section E.1.

4.4 Stable phase: stage 3

Data mixture

In the third and last stage of the stable phase (8T to 10T tokens, before annealing starts), we added the text-only English portion of InfiMM-WebMath with OWM, bringing the total proportion of math data to approximately 10%, as shown in Figure 2. For English web data, we revisited our ablation findings and adjusted the FineWeb-Edu to DCLM ratio to 40/60. For code, we replaced StarCoderData with Stack-Edu (Section 3.4). For languages with fewer than 4B tokens in Stack-Edu (TypeScript, Shell, Swift, Go, Rust and Ruby), we used their StarCoder2Data subsets. We also added Jupyter Notebooks from StarCoder2 (Lozhkov et al., 2024), which provides rich, contextual examples of code interleaved with explanations, enhancing the model’s reasoning around programming tasks.

Findings

While the integration of these new datasets brought improvements across multiple benchmarks, we observed a noticeable loss spike during this phase which remained even after rewinding training and skipping data associated with the spike (Chowdhery et al., 2023; Almazrouei et al., 2023). The exact cause remains undetermined but most evaluation metrics recovered by the end of the stage.

4.5 Decay phase: stage 4

Data mixture

The final stage consisted of decaying the learning rate linearly to 0 for 10% of the total training duration (from 10T to 11T tokens) (Hägele et al., 2024). Following Blakeney et al. (2024), we introduced our highest quality mathematical datasets, InfiWebMath-3+, and FineMath 4+. We additionally allocated 0.08% of the mixture to OWM and 0.02% to AugGSM8K (Li et al., 2024a), an augmented version of the GSM8K benchmark’s training set, which has become a common component of recent pretraining datasets (Achiam et al., 2023; Dubey et al., 2024; Ai2, 2024). Overall, mathematical content totaled 14% of the mixture. We expanded Stack-Edu to include additional programming languages not covered in stage 3, and set the dataset’s contribution to 24% of the mixture. We maintained the natural distribution across programming languages, with a higher allocation for Python. The remaining mixture consisted of English web data at 58% (maintaining the higher DCLM to FineWeb-Edu ratio) and Cosmopedia v2 (Allal et al., 2024) at 4%, which provides 30B tokens of high-quality synthetic textbooks, blog posts, and stories.

Findings

While all benchmark tasks show improvements after stage 4, we observe substantial gains in coding performance and, most notably, in math performance, validating our data mixture specifically targeting these domains.

4.6 Context Length extension

To support long-context applications, we followed standard practice (Gao et al., 2024) and extended the context length from 2k to 8k tokens, by taking an intermediate checkpoint from stage 4 (before the final 75 billion tokens of training) and continuing training with a different data mixture and a RoPE value of 130k. The mixture was adjusted to include 40% long-context documents (8k tokens or more) sourced from DCLM (10%), FineWeb-Edu (10%), and the books subset of Dolma (20%) (Soldaini et al., 2024), while the remaining 60% followed the stage 4 mixture. After this step, we obtain the final SmolLM2 base model.

4.7 Base model evaluation

We evaluate and compare the final base SmolLM2 model with existing state-of-the-art models of similar size, Qwen2.5-1.5B (Yang et al., 2024b) and Llama3.2-1B (AI@Meta, 2024a), on a wide range of benchmarks. Evaluations are conducted using lighteval and in a zero-shot setting unless otherwise specified.

Evaluation results in Table 4 show the strong performance of base SmolLM2, outperforming the Qwen2.5 base model on HellaSwag, and ARC. SmolLM2 also delivers strong performance on held-out benchmarks not monitored during training, such as MMLU-Pro (Wang et al., 2024c), TriviaQA (Joshi et al., 2017), and Natural Questions (NQ, Kwiatkowski et al., 2019). Notably, the model outperforms Qwen2.5-1.5B by nearly 6 percentage points on MMLU-Pro, further validating its generalization capabilities. On math and coding benchmarks, SmolLM2 demonstrates competitive performance. While it lags behind Qwen2.5-1.5B, SmolLM2 outperforms Llama3.2-1B on GSM8K, MATH and HumanEval. Importantly, we see next to no degradation in performance after Context Length Extension, while the HELMET (Yen et al., 2024) and Needle in the Haystack (NIAH) (Kamradt, 2024) results show strong performance – see Appendix G. These results highlight the effectiveness of our curated datasets, data mixtures, and training stages.

Table 4: Performance comparison of SmolLM2 and other 1-2B base models across benchmarks. SmolLM2 demonstrates competitive results highlighting its generalization capabilities.