2.3 Pre-training Workflow

2.3 Pre-training Workflow

2.3.1 Tokenizer Training

Tokenization Tips

• Digit splitting (avoid issues like 9.11 > 9.9)

• Control compression rate (compares the number of characters in a string to the number of tokens)
For example: "The quick brown fox jumps over the lazy dog.", where the total characters is 44 (including spaces and punctuation). Using a GPT-style tokenization, this sentence becomes 9 tokens: ['The', ' quick', ' brown', ' fox', ' jumps', ' over', ' the', ' lazy', ' dog.'] . Therefore,
• A higher compression rate means each token represents more characters, which is generally more efficient but being too high affects the model’s knowledge modeling ability.
• A lower compression rate might mean there's lots of punctuation, emojis, or short uncommon words (which take more tokens), resulting in lower encoding and decoding efficiency.

• Manually add domain-specific and remove sensitive tokens
If you're improving a model for your own business needs, you should add tokens that are relevant to your domain to increase compression rate for those specific texts. For example, in the medical domain, you might add tokens like "amoxicillin" or "penicillin", so the tokenizer won’t split them into meaningless tokens.

• Vocabulary should cover enough tokens
As for whether to include minority languages, that depends on business requirements.

• Reserved Token Count

Reserved tokens (also known as special tokens) are typically labeled as reserved_token, unused_token, or padded_token, all referring to the same concept. These tokens are not meant to appear in natural language data; instead, they are reserved for special usage during the post-training phase. For example, they may be used for task separators, role identifiers, special instructions in function calls, or commands for agents.

It's best to reserve enough of these tokens—typically between 100 and 1000. If the tokenizer you downloaded doesn't include enough reserved tokens, you can manually add more.

Vocabulary Expansion

2.3.2 Determine Model and Pre-Training Framework

2.3.3 Pre-training Strategy

Optimal batch_size

Warmup-Stable-Decay (WSD) Learning Rate Scheduler

• Warmup phase (denoted by W, the end of steps/iterations in the warmup phase),

• Stable training phase (denoted by T, the end of steps/iterations in the stable phase),

• Decay phase (denoted by S, the end of steps/iterations in the decay phase).

• Supports continuous training. The strategy is stable and robust, allowing the model to be trained smoothly without premature convergence or failure.

• Allows model extraction at any time. The learning rate remains reasonable throughout the training process, so the model can be used or fine-tuned at any training step.

• Outperforms cosine learning rate schedule. Compared to the commonly used cosine learning rate strategy, this method achieves better performance in practice.

• Explicit training phases enable flexible data strategies. With clearly defined warmup, stable, and decay phases, it is easier to apply different data augmentation or sampling strategies at each stage.

Pretraining Tricks

• Improve training efficiency: The training duration usually exceeds one month. If memory is sufficient, avoid using tensor_parallel, pipeline_parallel, or sequence_parallel. Similarly, if offload can be avoided, don't enable it; if recomputation can be avoided, then don't enable it either.

• Add instruction data: In the current view, the more instruction data, the better. This is reasonable because the more instructions are seen during pre-training, the easier domain adaptation in SFT (more similar distribution between pre-train and post-train data).

Training Setup

• Warm-up at the beginning, gradually increasing the learning rate to the target level

• Middle stage: use cos / cos_decay / constant / constant_decay, with a relatively large learning rate. Whether to apply decay depends; could refer to others’ technical reports or experiment on small-scale training

• Later stage: extend training sequence length, and change the frequency of RoPE to make the model adapt to longer texts

• Final stage: apply annealing for fine-tuning. Use high-quality datasets like IFT (Instruction Fine-Tuning) to strengthen the model’s instruction-following ability. After finishing this phase, move to testing or benchmarking

2.3.4 Training Monitoring

1. Monitor channel_loss. At the very least knowledge and code—these types of data loss should be observed separately.

2. Observe loss spike.
Currently, there is no clear evidence showing that the loss spike will cause irreversible model damage, but it’s better not to have a loss spike generally. Whether the spike is sudden or gradually increasing, it usually means there’s a data issue, either the data's loss is high or low. For example, when the data is all garbled (Mojibake), the loss is very high; when the data is all newline characters, the loss is very low.
💡
Solutions for Loss Spike
• Switch checkpoint to the step before the loss spike and retrain (from GLM130B technical report)
• Decrease the learning rate
• Multiply the shadow layers’ gradients with a scaler to reduce the update step
• Use training strategies such as WSD suggested by MiniCPM
• Use z-loss (used for final softmax logits in the Mesh Tensorflow codebase; also adopted by ST-MOE) to regularize and stabilize the softmax output

3. PPL (Perplexity) Monitoring: In general (e.g. the general corpus):
1. The PPL first increases then decreases;
2. The data with less ratio may show an increasing trend; added domain data (more ratio) is more likely to drop. We usually randomly sample 200 samples from a dataset (the same source) for PPL monitoring.
📌
Perplexity is a commonly used evaluation metric in natural language processing to measure the performance of a language model. A lower perplexity indicates better prediction performance of the model on the next word in a sentence.

2.3.5 Pre-training Scaling Law

1. First, it’s important to estimate computing resources.
💡
Proof
We define the model structure as follows:
• Decoder layers:
• Attention hidden dimension:
• Attention feedforward layer dimension: , usually
First, the number of parameters in the model, denoted as , (excluding embedding, norm, and bias) is calculated as:
Each transformer layer consists of self-attention and MLP components: 1. Self-Attention Parameters: , , , , each with dimensions , so total parameters are 2. MLP Parameters: , , so total parameters are Total Parameters Per Layer: Thus, total model parameters across all layers:
---
Next, we calculate forward pass computation using FLOPs (floating point operations). Note that For matrices , matrix multiplication FLOPs =
Assume decoder input:
• : batch size
• : sequence length
• : model dimension
Self-Attention Computation:
Input linear projection: ; FLOPs: Attention scores (): Value product (let’s ignore the mask and softmax here): Output projection ():
MLP Computation:
Feedforward Up (): Feedforward Down ():
Single Decoder Layer FLOPs:
Total Decoder Layer FLOPs (L layers):
---
Let’s assume the backward computation is twice that of forward, so the total FLOPs for forward and backward:
Per Token FLOPs:
Where the . so the total computation (for full dataset of D tokens):
Where the is the model size (parameters count) and the is the token count. QED. More details can be seen in Scaling Laws for Neural Language Models.

2. Next, let's estimate our budgets: how many GPUs * how many days. For example: 100 A800 training for one month. A single A800 can achieve approximately 312 TFLOPs/GPU without structural sparsity enabled, then the total compute is:

Note that:

3. So how do we determine the model size for training?
According to the formula, with a given compute budget, let’s say you want to train a 7B parameter model (7 billions parameters), the amount of data you can use is:

4. Llama3's scaling law
Given different budgets (from to FLOPs), the model size was adjusted (ranging from 40M to 16B). Then the following chart was drawn:



With the experimental results from the small model, the next step is to predict which model is most computationally optimal under large-scale computation. This paper still uses an empirical law to model:
Here, is the FLOPs budget for input, is the compute-optimal number of tokens under this budget, and and are parameters to be fitted. By using the optimal computation point (red points) from the previous experiment and redrawing the graph, we can fit this curve. Here, , .

5. Metrics Prediction
Use the ARC Challenge evaluation set. This is a relatively difficult multiple-choice evaluation set. A common evaluation method is to let the model calculate the PPL (Perplexity) of each option, and then see which option has the lowest PPL, which is considered the model’s final answer. The overall accuracy is then calculated.
The metrics predictions are divided into two steps:
• Use correct answers to fit the relationship between NLL loss and FLOPs (here use the small but optimal model in the previous experiment);
• Use sigmoid to fit the relationship between loss and accuracy (here the small model + Llama 2 is used for better prediction).



After getting those two curves, we know how good the model can be under a fixed budget. The LLaMA 3 authors discovered that this method is very accurate, and it can extrapolate across 4 orders of magnitude. On the LLaMA 3 405B model, the estimate is just slightly different from the real one.
📌
In models under 70B, "compute-optimal" is not very important, because at that time there's no fixed budget, and the given constraint is model size — the inference speed needs to be ensured. In larger-scale training scenarios, the constraint becomes compute budget, so "compute-optimal" becomes more important.
Perhaps in the future, as chips develop further, even 405B may become a "small model", and in that case, we might consider even more tokens, expanding the scale further.
Scaling Laws are about exploring what kind of model we can train within a fixed compute budget — they help us understand what we can get after investing this much money and emitting this much CO₂. Previous works mostly focused on predicting dev loss. LLaMA 3 authors let us feel its improved performance aligns with its predictions. When it comes to testing knowledge, this method looks fairly good, but for evaluating models like CoT or mathematics reasoning, better prediction methods are still needed.

2.3.6 Mixed Precision Training

1. Computation Preparation: Store an fp32 copy of the parameter, momentum, and variance (for the Adam optimizer). Then, duplicate the parameter and convert it to fp16, obtaining a fp16 copy of the parameter.
• The fp32 parameter acts as the "master weight". During model training, the updates done through optimizer.step() should apply to this master weight. When the model finishes training, this is the version of the weights we save. Our ultimate goal is to get highly accurate weights.
• The fp16 parameter acts as the "training weight", which is the weight used in actual training computations for the forward process.

2. FWD (Forward Pass): Use the fp16 parameter for forward computation. During this step, fp16 is also used for activation (the output from layers used in backpropagation). Especially note that if activation gradients need to be stored (e.g. when without recomputation), they can occupy a lot of memory (possibly exceeding the model itself). We use fp32 precision for the calculated loss to ensure numerical accuracy in backpropagation.

3. Loss Calculation: To prevent gradient underflow, scale the loss and get the fp32 scaled loss. We will get into this later.

4. BWD (Backward Pass): Perform backward computation using fp32 scaled loss. Since the loss was scaled in step 3, the gradients will also be scaled. To save memory, these scaled gradients are stored in fp16 format.

5. Unscaled gradients: Scaled gradients are stored in fp16 format, but when it's time to use the gradients to update the model weights, they must be converted to fp32 format.

6. Clip gradients: After converting the gradients to fp32, we can perform operations like clipping to further prevent gradient explosion or vanishing. More details will be explained later.

Why Mixed Precision Training?

1. Round-off Error:
A solution is to update the model weights after converting the gradients from fp16 (float16) to fp32.



Between , the fp16 representation has a fixed smallest segment of , i.e., the number larger than is , so the would be rounded off.

2. Gradient Underflow:
67% of the gradients have values smaller than (as a reminder, this is the lower bound of the range representable by fp16). In other words, if the entire training is done using fp16, in the later stages of training, gradient underflow will frequently occur, causing the whole training process to fail to proceed normally.

GPU Memory Usage

Loss Scale

• Static Loss Scaling:
Use a fixed scaling factor loss_scale to scale the loss. The process is as follows:
• Scale up phase:
In the backward phase, scale the loss value (in fp32) up by , compute the gradients with the scaled loss (in fp32), and store the gradients in fp16 precision.
• Scale down phase:
When we need to update the weights with gradients:
• First, check whether the scaled gradient contains overflow (inf/nan). (Although underflow issues are addressed by scaling up the loss, side effects of overflow may occur). If overflow is detected, skip this step of update.
• If no overflow is detected, unscale the fp16 gradients (i.e., divide the gradients by ), and convert them back to fp32 for the update.

• Dynamic Loss Scaling
• First, use a very large loss_scale, e.g., , then scale up the loss and compute gradients.
• Check whether the gradients after scaling up have overflow (inf/nan):
• If no overflow, convert gradients to fp32 and proceed with the normal update.
• If overflow is detected, skip this step of update and, at the same time, reduce the loss_scale by a factor F (default is 2).
• Once the model stabilizes and the gradient fluctuations become less intensive, we can try increasing the loss_scale. For example, every N iterations (default N is 2000), try increasing the loss_scale by a factor of F. If overflow is encountered after increasing, revert the loss_scale to the previous stable value, and repeat this strategy.

Clip gradients

• If , then

• If , then remains unchanged