Trainers

Pre-training

Large language models learn language representations/initialize model weights/learn probability distributions through unsupervised learning pre-training on a large general dataset. We expect that after pre-training, the model can handle large amounts of diverse datasets, and can then be fine-tuned through supervised learning to adapt to specific tasks.

When pre-training, please set stage to pt and ensure the dataset used conforms to the Pre-training Datasets format.

Below is a pre-training configuration example:

### model
model_name_or_path: meta-llama/Meta-Llama-3-8B-Instruct

### method
stage: pt
do_train: true
finetuning_type: lora
lora_target: all

### dataset
dataset: c4_demo
cutoff_len: 1024
max_samples: 1000
overwrite_cache: true
preprocessing_num_workers: 16

### output
output_dir: saves/llama3-8b/lora/sft
logging_steps: 10
save_steps: 500
plot_loss: true
overwrite_output_dir: true

### train
per_device_train_batch_size: 1
gradient_accumulation_steps: 8
learning_rate: 1.0e-4
num_train_epochs: 3.0
lr_scheduler_type: cosine
warmup_ratio: 0.1
bf16: true
ddp_timeout: 180000000

Post-training

After pre-training, the model’s parameters are initialized, and the model can understand semantics, grammar, and identify contextual relationships, performing well on general tasks. Although the model’s emergent zero-shot learning and few-shot learning capabilities enable it to complete specific tasks to a certain extent, prompts alone may not necessarily lead to satisfactory performance. Therefore, we need post-training to enable the model to perform well on specific tasks.

Supervised Fine-Tuning

Supervised Fine-Tuning is a method of training on a pre-trained model using a small-scale labeled dataset. Compared to pre-training a completely new model, supervised fine-tuning of an existing pre-trained model is a faster and more cost-effective approach.

When performing supervised fine-tuning, please set stage to sft. Below is a supervised fine-tuning configuration example:

...
stage: sft
finetuning_type: lora
...

RLHF

Since the data that language models learn from during supervised fine-tuning comes from the internet, the model may not follow user instructions well and may even output illegal or violent content. Therefore, we need to align model behavior with user needs. Through RLHF (Reinforcement Learning from Human Feedback) methods, we can further fine-tune the model through human feedback, enabling the model to follow user instructions better and more safely.

Reward model

However, obtaining real human data is very time-consuming and expensive. A natural idea is that we can train a reward model to replace humans in evaluating language model outputs. To train this reward model, we need to let the reward model learn human preferences, which is typically achieved by inputting human-annotated preference datasets. In preference datasets, data consists of three parts: input, good response, and bad response. The reward model is trained on preference datasets, allowing it to evaluate language model outputs more in line with human preferences.

When training a reward model, please set stage to rm, ensure the dataset used conforms to the Preference Dataset format, and specify the save path for the reward model. Here is an example:

...
stage: rm
dataset: dpo_en_demo
...
output_dir: saves/llama3-8b/lora/reward
...

PPO

After training the reward model, we can start the reinforcement learning phase of the model. Unlike supervised learning, reinforcement learning does not have labeled data. The language model takes prompts as input, and its output serves as input to the reward model. The reward model evaluates the language model’s output and returns the evaluation to the language model. Ensuring that both models run smoothly is a challenging task. One way to implement this is by using Proximal Policy Optimization (PPO). The main idea is that we want the language model’s output to receive as high ratings as possible from the reward model, while avoiding overly ‘aggressive’ changes to the language model. This approach allows the model to learn to align with human preferences without losing too much of its original problem-solving ability.

When using PPO for reinforcement learning, please set stage to ppo and specify the path to the reward model being used. Here is an example:

...
stage: ppo
reward_model: saves/llama3-8b/lora/reward
...

DPO

Since ensuring the good operation of both the language model and reward model simultaneously is challenging, one idea is that we can discard the reward model and train our language model directly based on human preferences, which greatly simplifies the training process.

When using DPO, please set stage to dpo, ensure the dataset used conforms to the Preference Datasets format, and set preference optimization related parameters. Here is an example:

...
### method
stage: dpo
pref_beta: 0.1
pref_loss: sigmoid  # choices: [sigmoid (dpo), orpo, simpo]
dataset: dpo_en_demo
...

KTO

KTO (Kahneman-Taversky Optimization) emerged to solve the problem of paired preference data being difficult to obtain. KTO uses a new loss function that only requires binary labeled data, meaning training only requires labeling whether responses are good or bad, and achieves similar or even better results than DPO.

When using KTO, please set stage to kto, set preference optimization related parameters, and use the KTO dataset.

Here is an example:

model_name_or_path: meta-llama/Meta-Llama-3-8B-Instruct
...
stage: kto
pref_beta: 0.1
...
dataset: kto_en_demo