juliannunezb/midtrain-mix-5b-gpt2

Midtrain-Mix-5B-GPT2是一个包含约50亿个token的高质量中间训练混合数据集，专门用于GPT-2词汇表模型的继续预训练/中间训练/退火阶段。数据集使用GPT-2 BPE分词器（词汇量50257）进行预分词，并以flat uint16二进制格式存储，可直接通过np.memmap加载。数据来源经过优化组合：70%来自FineWeb-Edu（经过教育价值筛选的网页内容），20%来自英文维基百科（密集的事实性文本），10%来自Open-Web-Math（数学和arXiv内容）。该数据集不包含代码，专注于提升模型在事实性、教育性和数学内容上的知识整合能力，适用于在通用网页数据预训练后进行的第二阶段训练，通过降低学习率和高质量数据混合来巩固模型知识。

语言：英语
许可证：Apache-2.0
任务类别：
- 文本生成
标签：
- 继续预训练
- 中间训练
- 退火
- GPT-2
- 已分词
- 事实性
- 教育性
规模类别：
- 10亿 < Token数 < 100亿
配置项：
- 配置名称：default
数据文件：
- 拆分集：train
路径：train.bin
- 拆分集：val
路径：val.bin

# Midtrain-Mix-5B-GPT2 数据集
**50亿Token**的高质量中间训练混合数据集，已使用**GPT-2 BPE分词器（GPT-2 BPE Tokenizer，词表大小50257）**完成预分词，打包为适配`np.memmap`的扁平`uint16`二进制文件。可直接接入任意适配GPT-2词表的预训练训练流程，用于已接触过通用网页语料的模型的**继续预训练/中间训练/退火**。

## 何为中间训练？其价值何在？
在网页数据（如FineWeb、RedPajama、DCLM等）上完成标准预训练阶段后，现代大语言模型（Large Language Model, LLM）流程会开展**第二轮更短的训练阶段：使用更低的学习率（Learning Rate, LR），基于更高质量的混合语料**。该技术在不同实验室有多种命名：

| 实验室 | 命名 |
|---|---|
| Meta（Llama-3项目） | 「退火阶段（annealing phase）」 |
| DeepSeek（V3/R1系列） | 「推理语料继续预训练」 |
| 艾伦人工智能研究所（Allen AI，OLMo项目） | 「第二阶段预训练」 |
| Hugging Face（SmolLM2项目） | 「中间训练（mid-training）」 |

该流程的核心逻辑一致：严格筛选事实性、教育性、数学类内容，将学习率降至原水平的1/10，训练时长约为原预训练预算的5%~20%。如此可让模型在不遗忘已有知识的前提下，将更多知识固化至权重参数中。

本数据集即为面向小型GPT-2词表模型的「优质混合语料」。

## 概览
- **总Token数**：约50亿（设计目标为在该数据集上完成1个训练轮次（epoch））
- **分词器**：`tiktoken.get_encoding("gpt2")`，词表大小50257，文档结束符（End of Text, EOT）ID为50256
- **存储格式**：扁平uint16二进制文件，可直接用于`np.memmap`内存映射
- **数据来源**（比例针对小型模型的知识密度需求优化）：
- 70% FineWeb-Edu（`HuggingFaceFW/fineweb-edu`的`sample-10BT`子集）：通过基于Llama-3-70B标注训练的分类器筛选出的教育类网页语料
- 20% 英文维基百科（Wikipedia (en)，`wikimedia/wikipedia`的`20231101.en`子集）：高密度事实性散文文本
- 10% Open-Web-Math（`open-web-math/open-web-math`）：数学、arXiv论文及LaTeX格式文本
- **本混合语料不含代码**。原因：在5000万至2亿参数的模型规模下，代码语料的利用率极低；而诸如[smoltalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk)这类监督微调（Supervised Fine-Tuning, SFT）语料已通过`apigen-80k`和`self-oss-instruct`包含代码数据。对于更大规模的模型，可另行添加代码语料源。

## 数据集统计信息
| 指标 | 训练集 | 验证集 |
| :--- | ---: | ---: |
| Token数 | ~49.7亿 | ~2500万 |
| 存储大小（uint16格式） | ~9.94 GB | ~50 MB |
| 验证集占比 | 0.005 | — |

（构建完成后，完整统计数据见`stats.json`文件）。
训练/验证集划分方式：基于每个文档以随机种子42进行抽样，验证集占比VAL_FRAC=0.005。验证集**完全匹配混合语料的分布**（即采用与训练集相同的抽样方式），并非仅从单一数据源中保留的数据。

## 数据格式规范
数据格式极为简洁，仅包含两个扁平文件：

train.bin uint16 拼接后的Token ID序列，为无分隔符的单一流数据
val.bin uint16 格式与训练集一致，为验证集Token序列

文档拼接时**不添加**EOT分隔符（训练流程的随机窗口采样器有时会采样到跨文档的窗口，这在中间训练场景中是可接受的，且通常有助于上下文融合）。

## 数据加载方式
### 标准`np.memmap`内存映射加载（推荐方式）
python
import numpy as np
train = np.memmap("train.bin", dtype=np.uint16, mode="r")
val = np.memmap("val.bin", dtype=np.uint16, mode="r")
print(f"train: {len(train)/1e9:.2f} B tokens, val: {len(val)/1e6:.1f} M tokens")

# 采样一个长度为BLOCK_SIZE的随机训练窗口（预训练风格）：
BLOCK_SIZE = 2048
ix = np.random.randint(0, len(train) - BLOCK_SIZE - 1)
window = np.asarray(train[ix:ix+BLOCK_SIZE+1], dtype=np.int64)
x = torch.from_numpy(window[:-1]) # 形状：(BLOCK_SIZE,)
y = torch.from_numpy(window[1:]) # 下一个Token预测目标

### 解码样本以验证正确性
python
import tiktoken
tok = tiktoken.get_encoding("gpt2")
sample_ids = list(train[1_000_000:1_000_500]) # 选取中间位置的500个Token
print(tok.decode(sample_ids))

## 中间训练使用方法
在完成通用预训练后，从预训练检查点（checkpoint, ckpt）恢复训练，并设置以下参数：
- **学习率（LR）**：为预训练峰值学习率的1/10。若预训练峰值学习率为1.5e-3，则中间训练阶段使用1.5e-4的学习率，并采用余弦退火至初始学习率的10%。
- **预热步数（WARMUP）**：设置为较短的步数（100~300步）——模型已完成预热。
- **总迭代步数（N_ITER）**：目标为在该数据集上完成1个训练轮次。当窗口长度为2048、批量大小为16时：

N_ITER = 50亿 / (2048 * 16) ≈ 152587 步

- **优化器**：使用全新的优化器（丢弃预训练的动量参数）。模型架构需与预训练阶段完全一致。
- **采样模式**：无重复洗牌模式（每个训练窗口在每个轮次中仅被采样一次）。

参考启动脚本示例（[`local50m`](https://github.com/)项目采用了该流程）：
bash
MODE=pretrain
DATA_DIR=/path/to/midtrain
RESUME_FROM=/path/to/pretrain_final.pt
RESUME_OPTIM=0 RESUME_STEP=0
LR=1.5e-4 MUON_LR=4.5e-3 MIN_LR_RATIO=0.1 WARMUP_STEPS=200
N_ITER=150000 CKPT_INTERVAL=10000 BLOCK_SIZE=2048 BATCH_SIZE=16
SAMPLE_MODE=shuffled
HF_CHECKPOINT_REPO=user/your-midtrain-model
python3 local50m_train.py

针对预训练充分的模型，预期训练行为如下：
- 在前10000步中，验证集损失会显著下降（知识固化阶段）。
- 约12万步左右，损失趋于平稳至余弦退火的最小值。
- 总训练时长：在RTX 5090显卡上约6小时，成本约3美元。

## 复现流程
bash
HF_TOKEN=... python3 prepare_midtrain.py # 仅本地运行
HF_TOKEN=... UPLOAD_TO=user/repo PRIVATE=0 python3 prepare_midtrain.py
SMOKE=1 python3 prepare_midtrain.py # 1000万Token冒烟测试
N_PROC=16 python3 prepare_midtrain.py # 并行分词

### 构建细节
- **构建时间**：2026年5月13日UTC，基于单台M1 Max设备（8个工作线程，处理50亿Token约需3小时）
- **构建脚本**：[`prepare_midtrain.py`](./prepare_midtrain.py)
- **上游数据集提交版本**：构建时的最新提交
- `HuggingFaceFW/fineweb-edu`（sample-10BT子集）
- `wikimedia/wikipedia`（20231101.en子集）
- `open-web-math/open-web-math`

### 已执行的合理性校验
- 所有Token ID均处于`[0, 50256]`区间（GPT-2 BPE分词器的有效范围）。
- 随机选取500Token窗口进行解码-编码往返测试，输出结果符合自然英文、数学文本或代码格式。
- 各数据源的Token占比与目标比例的误差不超过±2%。

## 流水线上下文
本数据集可作为**中间训练语料**嵌入完整的现代大语言模型流水线：

juliannunezb/mixed-pretrain-10b-gpt2 ← 通用预训练语料（98亿Token）
↓ 预训练
juliannunezb/midtrain-mix-5b-gpt2 ←━━ 本数据集（50亿Token，优质语料）
↓ 中间训练
juliannunezb/smoltalk-gpt2-sft ← 监督微调语料（100万对话数据）
↓ 监督微调
juliannunezb/ultrafeedback-gpt2-dpo ← 偏好配对数据（6万条）
↓ 直接偏好优化（Direct Preference Optimization, DPO）
[可选：在可验证任务上使用GRPO训练]

## 许可证
本数据集（分词后的Token流）采用**Apache-2.0许可证**，因其三个上游数据集均符合Apache-2.0或CC-BY-SA许可协议：
- FineWeb-Edu：Open Data Commons By Attribution 1.0
- 英文维基百科：CC-BY-SA-3.0 + GFDL
- Open-Web-Math：Open Data Commons By Attribution 1.0

完整来源信息请参阅各上游数据集的卡片页面。

## 引用方式
若使用本数据集，请引用以下上游数据源：
bibtex
@misc{lozhkov2024fineweb-edu,
author = {Lozhkov, Anton and Ben Allal, Loubna and von Werra, Leandro and Wolf, Thomas},
title = {FineWeb-Edu: the Finest Collection of Educational Content the Web Has to Offer},
year = {2024},
url = {https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu}
}
@misc{paster2023openwebmath,
title = {OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text},
author = {Paster, Keiran and Dos Santos, Marco and Azerbayev, Zhangir and Ba, Jimmy},
year = {2023},
eprint = {2310.06786},
archivePrefix = {arXiv}
}

若引用本数据集的打包版本，请引用：
bibtex
@misc{midtrain_mix_5b_gpt2_2026,
title = {Midtrain-Mix-5B-GPT2: a premium mid-training mix for small GPT-2-vocab models},
author = {Julián Núñez},
year = {2026},
publisher = {Hugging Face},
url = {https://huggingface.co/datasets/juliannunezb/midtrain-mix-5b-gpt2}
}

## 已知局限性
- **仅支持英语**：三个上游数据源均以英语为核心。
- **仅适配GPT-2 BPE分词器**：若目标模型使用Llama或NeoX的BPE分词器，需基于原始数据源重新进行分词。
- **50亿Token对于小型模型而言数据量充足，但对于大型模型则不足**：最适合5000万至5亿参数的模型；对于10亿参数以上的模型，需约500亿同质量的语料。
- **网页筛选存在缺陷**：FineWeb-Edu的分类器偏向于正式/学术性文本，对工程、教程、爱好者类内容的采样不足。若需更广泛的覆盖范围，可按较低权重混入未筛选的FineWeb语料。
- **本混合语料不含领域特定代码**：若需支持代码的模型，请使用单独的代码混合语料（如StarCoderData、The Stack）。