MATES - GitLite

KO | EN

gitlite — search

#javascript #typescript #llm #golang #android #machine-learning #deep-learning #docker #macos #nodejs #awesome #cli #pytorch #ios

MATES

(★ 79)

Official repository for MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models [NeurIPS 2024]

File Explorer

Download Latest Version (.zip)

- avatar.png
- MATES.png
- eval.sh
- predict_data_influence.sh
- pretrain.sh
- probe_oracle_data_influence.sh
- train_data_influence_model.sh
- - __init__.py
  - config.py
  - gen.py
  - model.py
  - packed_dataset.py
  - rmsnorm.py
  - tokenizer.py
  - utils.py
- - lm_eval_harness.py
  - pretrain.py
- - modeling_data_influence_model.py
  - predict_data_influence.py
  - prepare_lambada.py
  - probe_oracle_data_influence.py
  - select_data.py
  - train_data_influence_model.py
.gitignore
LICENSE.txt
README.md
requirements.txt

# Installation Guide

1. Get the code

git clone https://github.com/cxcscmu/MATES

Downloads the entire project code from GitHub to your computer.

cd MATES

Moves into the project folder you just downloaded.

2. Python

Easy Recommended

Prerequisites

Git Needed to download the project code from GitHub.
Python 3 On Windows, be sure to check 'Add Python to PATH' during installation.

pip install -r requirements.txt

Installs the Python libraries listed in requirements.txt (or similar).

✅ If it runs without errors and prints output in the terminal, it worked.

Pulled directly from this repo's README.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# [NeurIPS 2024] MATES<img src="assets/avatar.png" alt="drawing" style="height: 1em;">: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models

<p align="center"><a href='https://huggingface.co/yuzc19/pythia-410m-mates'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Main Model-blue'> <a href='https://huggingface.co/yuzc19/bert-base-uncased-data-influence-model-lambada'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Data Influence Model-blue'></p>

This is the official repository for [MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models](https://arxiv.org/pdf/2406.06046). The implementation is mainly based on [LitGPT](https://github.com/Lightning-AI/litgpt), which is easy to begin with, use, and modify.

<br>
<p align="center">
<img src="assets/MATES.png" width="600">
</p>
<br>

## Quick Links

- [1 Environment](#1-environment)
- [2 Dataset](#2-dataset)
- [3 Experiments](#3-experiments)
  - [3.1 Pretraining](#31-pretraining)
  - [3.2 Data Selection](#32-data-selection)
  - [3.3 Evaluation](#33-evaluation)
- [4 Citation](#4-citation)

## 1 Environment

**Python version**

The code is tested on Python 3.9.17.

**Install basic dependencies**

```bash
pip install -r requirements.txt
```

## 2 Dataset

We use a tokenized version of the [C4 dataset](https://huggingface.co/datasets/loganengstrom/dsdm-candidate-c4) in our code. Please ensure your disk has at least 500 GB of storage for this dataset. To get the training data for the initial warmup 10k steps, please run:

```bash
python src/select_data/select_data.py
```

- The selected data will be saved in `data/c4/pythia-410m/random/0`.

For preprocessing our reference task LAMBADA, please run:

```bash
python src/select_data/prepare_lambada.py
```

- The processed data will be saved in `data/lambada_openai`.

## 3 Experiments

Our main experiments use 8 GPUs for parallelization.

### 3.1 Pretraining

Our pretraining is run stage by stage to facilitate the model-aware data selection. Each stage consists of 10k steps. For instance, in the initial warmup 10k steps, you can run:

```bash
model_name=pythia-410m \
method=random \
ckpt=0 \
decay=false \
bash scripts/pretrain.sh
```

- `ckpt=0` denotes we are training from scratch.

To resume the pretraining from previous steps (e.g., 10k), you can run:

```bash
model_name=pythia-410m \
method=random \
ckpt=40000 \
decay=false \
bash scripts/pretrain.sh
```

- `ckpt=40000` denotes our gradient accumulation step is 4.
- `method=random` is the random data selection. You can replace it with `mates` for MATES after the first 10k steps.

### 3.2 Data Selection

After the first 10k steps, we can start the MATES data selection process every 10k steps. One data selection process consists of four steps:

1️⃣ Get oracle data influence:

```bash
model_name=pythia-410m \
method=random \
ckpt=40000 \
bash scripts/probe_oracle_data_influence.sh
```

- For the 10k checkpoint, `method=random`, but for the following, `method=mates`.

2️⃣ Train data influence model:

```bash
model_name=pythia-410m \
ckpt=40000 \
bash scripts/train_data_influence_model.sh
```

3️⃣ Predict data influence:

```bash
model_name=pythia-410m \
ckpt=40000 \
bash scripts/predict_data_influence.sh
```

4️⃣ Select the training data for the next 10k steps:

```bash
python src/select_data/select_data.py --model_name pythia-410m --method mates --ckpt 40000
```

- The selected data will be saved in `data/c4/pythia-410m/mates/40000`.

### 3.3 Evaluation

1️⃣ It is advised to run the evaluation after the decay stage for intermediate checkpoints for better stability.

```bash
model_name=pythia-410m \
method=mates \
ckpt=80000 \
decay=true \
bash scripts/pretrain.sh
```

2️⃣ We provide a simple evaluation example here and you can modify the parameters based on your needs.

```bash
model_name=pythia-410m \
method=mates \
ckpt=80800 \
bash scripts/eval.sh
```

- After running the evaluation script, you can find the results in the `results/c4/$model/$method/iter-$ckpt-ckpt/results.json`.

## 4 Citation

Please cite our paper if you use MATES in your work:

```bibtex
@inproceedings{yu2024mates,
   title={MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models},
   author={Yu, Zichun and Das, Spandan and Xiong, Chenyan},
   booktitle={NeurIPS},
   year={2024}
}
```