PreSelect-100B
收藏Predictive Data Selection: The Data That Predicts Is the Data That Teaches
简介
本研究旨在通过高效方式选择预训练数据，直接估计数据在预训练中的贡献。基于模型在某些文本上的压缩效率（即归一化损失）与其下游性能强相关的观察，我们假设那些模型损失预测下游能力的文本也有效地促进了学习。为此，我们引入了基于数据预测强度（PreSelect）的数据选择方法，这是一个轻量级且高效的数据选择方法，仅需要训练和部署一个基于fastText的评分器。通过在1B和3B参数模型上进行综合实验，我们证明了在30B标记上训练的PreSelect性能超过了在300B标记上训练的基线，实现了计算需求的10倍减少。此外，PreSelect显著优于其他竞争性数据选择基线，如DCLM和FineWeb-Edu。
发布资源
我们发布了训练好的fastText分类器和在Huggingface上经过100B标记过滤的高质量数据集，供直接使用。
| 名称 | 类型 | Huggingface链接 |
|---|---|---|
| preselect-fasttext-calssifier | 模型 | 🤗Huggingface |
| preselect-100B | 数据集 | 🤗Huggingface |
环境设置
我们提供了一个Dockerfile，其中包含了过滤、训练和评估的环境。
bash docker build -t preselect:latest . docker run --gpus all --network host -it --shm-size=20g --privileged preselect:latest
之后，你需要准备你的预训练语料库（例如下载commoncrawl子集）。我们提供了一个示例来下载DCLM的Refinedweb。注意这需要你事先设置aws服务。
bash cd data_processing/data/clean_pool pythondownload.py python unzip.py
你还可以准备自己的数据。
过滤
直接使用
如果你想直接使用我们训练好的fastText，你可以从Huggingface下载并运行以下代码：
python import os import argparse
from pathlib import Path
parser = argparse.ArgumentParser("Filter") parser.add_argument("--input_path", type=str, help="input path name") parser.add_argument("--output_path", type=str, help="output path name")
args = parser.parse_args() from datatrove.executor import LocalPipelineExecutor from datatrove.pipeline.filters import FastTextClassifierFilter from datatrove.pipeline.readers import ParquetReader, JsonlReader from datatrove.pipeline.writers.jsonl import JsonlWriter Path(f"{args.output_path}").mkdir(parents=True, exist_ok=True)
dist_executor = LocalPipelineExecutor( skip_completed=True, pipeline=[ JsonlReader(f"{args.input_path}", text_key="text", default_metadata={}), FastTextClassifierFilter(f"PreSelect-classifier.bin", keep_labels=[("1", 0.5)]), JsonlWriter(f"{args.output_path}", compression=None) ], tasks=100, ) dist_executor.run()
训练自己的FastText
第一步：对子集进行BPC计算
首先，挑选一个小子集并为每个模型的每个示例计算bpc。
bash cd data_processing/bpc python -u main.py --model_name {MODEL_NAME} --block_size 1900 --stride 512 --batch_size 1
第二步：训练Fasttext
然后，你可以使用第1步计算的数据训练fasttext。
bash cd data_processing/fasttext python train_fasttext.py
第三步：过滤
最后，你可以使用fasttext过滤你的大型语料库。提供的脚本在单CPU机器上工作，但可以轻松扩展到多机器过滤。
bash bash pipelie.sh {FASTTEXT_NAME} filter NO NO NO NO 0 NO 1 0.1
训练
如果你使用单个节点（例如8个GPU）进行训练，可以使用以下命令：
bash bash pipeline.sh {FASTTEXT_NAME} NO tokenize train convert NO 0 NO 1 0.1 {HOME_PATH} 1 {TRAINING_STEPS}
如果你使用多节点（例如8个GPU * 4节点）进行训练，可以使用以下命令：
bash bash pipeline_multi_node.sh {FASTTEXT_NAME} NO tokenize train convert NO {MAIN_NODE_ADDRESS} NO 1 0.1 {HOME_PATH} {N_NODE} {TRAINING_STEPS}
更多信息，请参考管道脚本。
评估
你可以参考Opencompass和LM-Evaluation-Harness来设置训练检查点的评估，以适应你的需求。
引用
如果你发现这项工作有帮助，请友好地引用如下：
@article{shum2025predictivedataselectiondata, title={Predictive Data Selection: The Data That Predicts Is the Data That Teaches}, author={Kashun Shum and Yuzhen Huang and Hongjian Zou and Ding Qi and Yixuan Liao and Xiaoxin Chen and Qian Liu and Junxian He}, journal={arXiv preprint arXiv:2503.00808}, year={2025}, eprint={2503.00808}, }
致谢
感谢以下开源项目的开源，其中一些代码在本项目中进行了引用和修改：
