Data Efficacy for Language Model Training

The paper describing this algorithm is Data Efficacy for Language Model Training: PDF

Forked from the Data Selection via Optimal Control project
GitHub: https://github.com/microsoft/LMOps/tree/main/data_selection

Overview of the Training Framework

1. Solve data quality scores based on a small model (160M), small total steps (100), SGD, and small proxy data (163,840 samples).
2. Fit the scores on the small proxy data with a 125M fairseq dense model (data scorer).
3. Annotate scores for large data (entire CC) with the data scorer.
4. Filter CC with the scores.
5. Pre-train the model.

Pre-Trained Models

• Models Trained on Redpajama CC (Conventional Pre-Training, Baselines)
• Models Trained PDS-Selected Data

Details of the Pipeline & How to run

The following pipeline can also be found in scripts/pipeline.sh.

0. Environment

conda create -n data_selection python=3.10
conda activate data_selection
bash install.sh
export BASE_PATH=/PATH/TO/THIS/PROJECT
# or export BASE_PATH=$PWD

1. Prepare Dataset

1.1 Tokenize CC for Pre-training

First download Redpajama CC. Then run tokenization.

bash $BASE_PATH/scripts/tools/process_cc.sh $BASE_PATH

1.2 Convert tokenized CC for data scorer inference if the data scorer and pre-training use different tokenizer

bash $BASE_PATH/scripts/tools/convert_tokenization.sh $BASE_PATH

1.3 Sample proxy data

bash $BASE_PATH/scripts/tools/sample_proxy_data.sh $BASE_PATH

2. Pre-train a small model

bash $BASE_PATH/scripts/pretrain/160M_bsl.sh $BASE_PATH

3. Solve Gamma

3.1 Prepare data for Algo

First download lima. Then run tokenization.

bash $BASE_PATH/scripts/tools/process_data/lima.sh $BASE_PATH

3.2 Run Algorithm

bash $BASE_PATH/scripts/pmp_solver/proxy_algo_160M.sh $BASE_PATH