Cosmopedia：如何为预训练创建大规模合成数据

发布于 2024 年 3 月 20 日

在 GitHub 上更新

卢布纳·本·阿拉尔 (Loubna Ben Allal)

在这篇博客文章中，我们概述了生成数十亿个 token 的合成数据集以复现 Phi-1.5 所面临的挑战和解决方案，从而促成了 Cosmopedia 的创建。合成数据已成为机器学习中的核心议题。它指的是由大型语言模型（LLM）等人工生成的数据，旨在模仿真实世界的数据。

传统上，为监督微调和指令微调创建数据集需要聘请人工标注员，这是一个成本高昂且耗时的过程。这种做法需要大量的资源，将此类数据集的开发限制在领域内的少数几个主要参与者。然而，最近情况发生了变化。我们看到了数百个高质量的合成微调数据集的开发，主要使用 GPT-3.5 和 GPT-4。社区也通过大量出版物支持了这一发展，这些出版物为各个领域的流程提供了指导，并解决了相关挑战 [1][2][3][4][5]。

number of datasets with synthetic tag
图 1. Hugging Face hub 上带有合成标签的数据集。

然而，这并非另一篇关于生成合成指令微调数据集的博客文章，社区已对此主题进行了广泛探索。我们专注于将样本数量从**数千**扩展到**数百万**，这些样本可用于**从头开始预训练 LLM**。这带来了一系列独特的挑战。

为什么选择 Cosmopedia？

微软通过其 Phi 系列模型 [6][7][8] 推动了这一领域的发展，这些模型主要在合成数据上进行训练。它们超越了在网络数据集上训练时间更长的大型模型。Phi-2 在过去一个月内被下载了超过 61.7 万次，是 Hugging Face hub 上最受欢迎的前 20 个模型之一。

尽管 Phi 模型的技术报告，例如《教科书是你所需要的一切》论文，揭示了这些模型卓越的性能和创建过程，但它们却省略了合成训练数据集策展的大量细节。此外，这些数据集本身并未发布。这引发了爱好者和怀疑者之间的争论。一些人称赞这些模型的能力，而批评者则认为它们可能只是过度拟合基准；其中一些人甚至将预训练模型在合成数据上的方法标记为“垃圾进，垃圾出”。然而，完全控制数据生成过程并复制 Phi 模型的高性能这一想法仍然引人入胜，值得探索。

这就是开发 Cosmopedia 的动力，它旨在复现用于 Phi-1.5 的训练数据。在这篇文章中，我们将分享我们的初步发现，并讨论改进当前数据集的一些计划。我们将深入探讨创建数据集的方法，深入了解提示词精选方法和技术栈。Cosmopedia 是完全开放的：我们发布了端到端管道的代码、数据集，以及一个基于它训练的 1B 模型，名为 cosmo-1b。这使得社区能够复现结果并在此基础上进行构建。

Cosmopedia 创建幕后

除了缺乏关于 Phi 数据集创建的信息外，另一个缺点是它们使用专有模型来生成数据。为了解决这些不足，我们推出了 Cosmopedia，这是一个由 Mixtral-8x7B-Instruct-v0.1 生成的合成教科书、博客文章、故事、帖子和 WikiHow 文章的数据集。它包含超过 3000 万个文件和 250 亿个 token，使其成为迄今为止最大的开放合成数据集。

注意：如果您期待在数百个 H100 GPU 上部署大规模生成任务的故事，那么实际上，Cosmopedia 的大部分时间都花在了细致的提示工程上。

提示词精选

生成合成数据看似简单，但大规模扩展时，保持多样性（这对最佳性能至关重要）变得极具挑战性。因此，精心策划涵盖广泛主题的多种提示词至关重要，并最大限度地减少重复输出，因为我们不想将计算资源花费在生成数十亿本教科书上，结果却因它们彼此高度相似而不得不丢弃大部分。在数百个 GPU 上启动生成之前，我们花费了大量时间使用 HuggingChat 等工具迭代提示词。在本节中，我们将介绍为 Cosmopedia 创建超过 3000 万个提示词的过程，这些提示词涵盖数百个主题，且重复内容少于 1%。

Cosmopedia 旨在生成大量高质量的合成数据，并具有广泛的主题覆盖。根据 Phi-1.5 技术报告，作者精选了 2 万个主题来生成 200 亿个 token 的合成教科书，同时使用网络数据集中的样本来增加多样性，报告中指出：

我们精心选择了 2 万个主题来启动这个新合成数据的生成。在我们的生成提示词中，我们使用网络数据集中的样本来增加多样性。

假设平均文件长度为 1000 个 token，这表明需要使用大约 2000 万个不同的提示。然而，结合主题和网络样本以增加多样性的方法仍不清楚。

我们结合两种方法来构建 Cosmopedia 的提示词：基于精选来源进行条件化和基于网络数据进行条件化。我们将用于条件化的数据来源称为“种子数据”。

piecharts of data sources
图 2. 构建 Cosmopedia 提示词的数据源分布（左图）和精选源类别内部的来源分布（右图）。

精选来源

我们使用来自斯坦福大学课程、可汗学院、OpenStax 和 WikiHow 等知名教育资源的专题。这些资源涵盖了 LLM 学习的许多有价值的专题。例如，我们提取了各种斯坦福大学课程的大纲，并构建了提示词，要求模型为这些课程中的各个单元生成教科书。图 3 展示了此类提示词的一个示例。

尽管这种方法能产出高质量内容，但其主要限制是可伸缩性。我们受限于每个来源的资源数量和可用专题。例如，我们只能从 OpenStax 中提取 16,000 个独特的单元，从斯坦福大学中提取 250,000 个单元。考虑到我们生成 200 亿个 token 的目标，我们至少需要 2000 万个提示！

利用受众和风格多样性

增加生成样本多样性的一种策略是利用受众和风格的多样性：通过改变目标受众（例如，幼儿与大学生）和生成风格（例如，学术教科书与博客文章），同一个主题可以多次重新利用。然而，我们发现仅仅将提示词从“为大学生编写一本关于‘为什么要进入太空？’的详细教科书单元”修改为“编写一篇关于‘为什么要进入太空？’的详细博客文章”或“为幼儿编写一本关于‘为什么要进入太空？’的教科书”不足以防止高重复内容率。为了缓解这个问题，我们强调了受众和风格的变化，提供了关于格式和内容应如何不同的具体说明。

图 3 展示了我们如何根据同一主题为不同受众调整提示词。

comparison of prompts
图 3. 为幼儿、专业人士和研究人员以及高中生生成相同教科书的提示词。

通过针对四种不同受众（幼儿、高中生、大学生、研究人员）并利用三种生成风格（教科书、博客文章、WikiHow 文章），我们可以获得多达 12 倍的提示。然而，我们可能希望包含这些资源中未涵盖的其他主题，并且这些资源的小批量仍然限制了这种方法，距离我们目标中的 2000 万+ 提示还很远。这时网络数据就派上用场了；如果我们生成涵盖所有网络主题的教科书会怎样？在下一节中，我们将解释我们如何选择主题并使用网络数据构建数百万个提示。

网络数据

使用网络数据构建提示词被证明是最具扩展性的方法，它贡献了 Cosmopedia 中超过 80% 的提示词。我们利用 RefinedWeb 等数据集，将数百万个网络样本聚类成 145 个簇，并通过提供 10 个随机样本的摘要并要求 Mixtral 找出它们的共同主题来识别每个簇的主题。关于此聚类的更多详细信息可在技术栈部分找到。

我们检查了这些群集，并排除了任何被认为教育价值低的内容。已删除内容的示例包括露骨的成人内容、名人八卦和讣告。保留和删除的 112 个主题的完整列表可以在这里找到。

然后，我们通过指示模型生成与网络样本相关的教科书来构建提示，该教科书基于聚类属于其主题的范围。图 4 提供了一个基于网络的提示示例。为了增强多样性并弥补主题标签中可能存在的任何不完整性，我们仅在 50% 的时间里根据主题来调整提示，并更改受众和生成风格，如上一节所述。我们最终使用这种方法构建了 2300 万个提示。图 5 显示了 Cosmopedia 中种子数据、生成格式和目标受众的最终分布。

web prompt
图 4. 网络提取和相关提示词示例。

histogram
图 5. Cosmopedia 数据集中种子数据、生成格式和目标受众的分布。

除了随机网络文件外，我们还使用了来自 AutoMathText 的样本，这是一个精心策划的数学文本数据集，旨在包含更多科学内容。

指令数据集和故事

在我们对使用生成的教科书训练的模型进行初步评估时，我们观察到这些模型缺乏小学教育中常见的常识和基本知识。为了解决这个问题，我们使用来自 UltraChat 和 OpenHermes2.5 指令微调数据集的文本作为提示词的种子数据，创建了包含日常知识和基本常识的故事。这些数据集涵盖了广泛的主题。例如，我们从 UltraChat 中使用了“关于世界的问答”子集，它涵盖了 30 个关于世界的元概念。对于 OpenHermes2.5，另一个多样化且高质量的指令微调数据集，我们省略了不适合讲故事的来源和类别，例如用于编程的 glaive-code-assist 和用于高级化学的 camelai。图 6 显示了我们用于生成这些故事的提示词示例。

stories prompts
图 6. 用于从 UltraChat 和 OpenHermes 样本为幼儿、普通受众和 Reddit 论坛生成故事的提示词。

这就是我们为构建 3000 多万个多样化提示而进行的提示工程故事的结局，这些提示提供了重复内容极少的内容。下图显示了 Cosmopedia 中存在的聚类，该分布类似于网络数据中的聚类。您还可以在Nomic 这里找到一个可点击的地图。

clusters
图 7. Cosmopedia 的聚类，使用 Mixtral 标注。

您可以使用数据集查看器自行查看数据集

dataset viewer
图 8. Cosmopedia 数据集查看器。

技术栈

我们发布了所有用于构建 Cosmopedia 的代码，网址为：https://github.com/huggingface/cosmopedia

在本节中，我们将重点介绍用于文本聚类、大规模文本生成以及训练 cosmo-1b 模型的技术栈。

主题聚类

我们使用 text-clustering 仓库来实现 Cosmopedia 提示词中使用的网络数据的主题聚类。下图展示了查找和标注聚类的管道。我们还要求 Mixtral 在标注步骤中为聚类给出 10 分制下的教育分数；这有助于我们在主题检查步骤中。您可以在此演示中找到网络聚类及其分数的演示。

text-clustering
图 9. 文本聚类的管道。

大规模教科书生成

我们利用 llm-swarm 库，使用 Mixtral-8x7B-Instruct-v0.1 生成了 250 亿个合成 token。这是一个可扩展的合成数据生成工具，可以使用本地 LLM 或 Hugging Face Hub 上的推理端点。它支持 TGI 和 vLLM 推理库。我们在 Hugging Face 科学集群的 H100 GPU 上本地部署了 Mixtral-8x7B，并使用 TGI。生成 Cosmopedia 的总计算时间超过 1 万 GPU 小时。

以下是一个示例，展示了如何在 Slurm 集群上使用 2 个 TGI 实例，对 10 万个 Cosmopedia 提示词使用 Mixtral 进行生成：

# clone the repo and follow installation requirements 
cd llm-swarm
python ./examples/textbooks/generate_synthetic_textbooks.py \
    --model mistralai/Mixtral-8x7B-Instruct-v0.1 \
    --instances 2 \
    --prompts_dataset "HuggingFaceTB/cosmopedia-100k" \
    --prompt_column prompt \
    --max_samples -1 \
    --checkpoint_path "./tests_data" \
    --repo_id "HuggingFaceTB/generations_cosmopedia_100k" \
    --checkpoint_interval 500

您甚至可以使用 `wandb` 跟踪生成过程，以监控吞吐量和生成的 token 数量。

text-clustering
图 10. llm-swarm 运行的 Wandb 图表。

注意：我们最初使用 HuggingChat 进行提示词迭代。然后，我们使用 `llm-swarm` 为每个提示词生成了数百个样本，以发现异常模式。例如，模型在教科书中使用非常相似的开场白，并且经常以相同的短语（如“曾几何时”和“太阳低悬在天空中”）开始故事。明确要求模型避免这些开场白并进行创新解决了这个问题；它们仍然被使用，但频率较低。

基准去污染

鉴于我们生成的是合成数据，种子样本或模型的训练数据中可能存在基准污染。为了解决这个问题，我们实施了去污染管道，以确保我们的数据集不包含任何来自测试基准的样本。

与 Phi-1 类似，我们使用 10-gram 重叠来识别潜在受污染的样本。在检索到候选样本后，我们使用 `difflib.SequenceMatcher` 将数据集样本与基准样本进行比较。如果 `len(matched_substrings)` 与 `len(benchmark_sample)` 的比率超过 0.5，则丢弃该样本。此去污染过程应用于所有使用 Cosmo-1B 模型评估的基准，包括 MMLU、HellaSwag、PIQA、SIQA、Winogrande、OpenBookQA、ARC-Easy 和 ARC-Challenge。

我们报告从每个数据集分割中移除的受污染样本数量，以及它们所对应的唯一基准样本数量（括号内）。

数据集组	ARC	BoolQ	HellaSwag	PIQA
网络数据 + 斯坦福 + OpenStax	49 (16)	386 (41)	6 (5)	5 (3)
auto_math_text + 可汗学院	17 (6)	34 (7)	1 (1)	0 (0)
故事	53 (32)	27 (21)	3 (3)	6 (4)

我们发现 MMLU、OpenBookQA 和 WinoGrande 的受污染样本少于 4 个。

训练栈

我们使用 Llama2 架构在 Cosmopedia 上训练了一个 1B LLM，以评估其质量：https://huggingface.co/HuggingFaceTB/cosmo-1b。

我们使用 datatrove 库进行数据去重和分词，使用 nanotron 进行模型训练，并使用 lighteval 进行评估。

该模型在 ARC-easy、ARC-challenge、OpenBookQA 和 MMLU 上的表现优于 TinyLlama 1.1B，并且在 ARC-challenge 和 OpenBookQA 上与 Qwen-1.5-1B 相当。然而，我们注意到与 Phi-1.5 相比存在一些性能差距，这表明更好的合成生成质量，这可能与用于生成的 LLM、主题覆盖或提示有关。

evaluation results
图 10. Cosmo-1B 的评估结果。

结论与后续步骤

在这篇博客文章中，我们概述了创建 Cosmopedia 的方法，这是一个旨在用于模型预训练的大型合成数据集，目标是复刻 Phi 数据集。我们强调了精心设计提示词以涵盖广泛主题、确保生成多样化内容的重要性。此外，我们分享并开源了我们的技术栈，它允许在数百个 GPU 上扩展生成过程。

然而，这只是 Cosmopedia 的初始版本，我们正在积极致力于提高生成内容的质量。生成内容的准确性和可靠性很大程度上取决于生成中使用的模型。具体来说，Mixtral 有时可能会出现幻觉并产生不正确的信息，例如在 AutoMathText 和 KhanAcademy 子集中的历史事实或数学推理方面。缓解幻觉问题的一种策略是使用检索增强生成（RAG）。这涉及检索与种子样本相关的信息（例如来自维基百科），并将其整合到上下文中。幻觉测量方法也有助于评估哪些主题或领域受其影响最大 [9]。比较 Mixtral 与其他开放模型的生成结果也很有趣。

合成数据的潜力巨大，我们热切期待社区将在 Cosmopedia 的基础上构建出什么。