训练大型语言模型通常涉及从网络爬取数据(例如 Common Crawl)、书籍、代码库和专门语料库等来源收集多样化的文本数据。在此过程中,一个主要考虑因素是确定每种数据来源在训练中的比例贡献。简单地将所有可用数据拼接并均匀打乱,看起来可能是一个直接的方法,但很少能产生最佳结果。这是因为不同的数据来源在质量、相关性和语言风格方面本身就存在差异。来源加权提供了对训练数据组合构成精确的控制,使得模型能够优先学习某些类型的文本而非其他。主要思路是为每个数据来源分配一个特定的权重或概率。在训练期间,当组建一个数据批次时,样本会根据这些预设权重从不同来源中抽样。与权重较低的来源相比,被分配较高权重的来源将随着时间为训练过程提供更多样本。定义抽样概率假设你有 $k$ 个不同的数据来源,$S_1, S_2, \dots, S_k$。每个来源 $S_i$ 包含 $N_i$ 个文档或样本。我们希望定义在任意给定步骤从来源 $S_i$ 抽样一个样本的概率 $p_i$。一个常见做法是使这个概率与来源的大小和分配的权重 $w_i$ 成比例:$$ p_i \propto w_i \cdot N_i $$或者,或许更直接的控制方式是,你可以直接为每个来源定义所需的比例(或概率)$p_i$,使得 $\sum_{i=1}^{k} p_i = 1$。这时,权重 $w_i$ 就简单地成为这些目标比例 $p_i$。例如,你可能会决定采用以下组合方式:网络文本 (来源 1): $p_1 = 0.60$ (60%)书籍 (来源 2): $p_2 = 0.15$ (15%)代码 (来源 3): $p_3 = 0.15$ (15%)维基百科 (来源 4): $p_4 = 0.10$ (10%)这意味着,平均而言,在每个训练周期(或大量步骤中),模型看到的样本中有 60% 将来自网络文本数据集,15% 来自书籍,以此类推。{"data": [{"type": "bar", "x": ["网络文本", "书籍", "代码", "维基百科"], "y": [60, 15, 15, 10], "marker": {"color": ["#339af0", "#82c91e", "#f76707", "#ae3ec9"]}}], "layout": {"title": "数据来源比例示例 (%)", "xaxis": {"title": "数据来源"}, "yaxis": {"title": "百分比"}, "margin": {"l": 40, "r": 20, "t": 60, "b": 40}, "height": 300}}大型语言模型预训练期间从不同数据来源抽样的比例示例。设定权重的方法和考量选择合适的权重通常受多种因素影响,需要仔细考虑并进行多次试验:数据质量: 维基百科或同行评审文章等高质量、精心整理的来源,相对于其大小,可能会比嘈杂的网络爬取数据获得更高的权重。这能指导模型学习更多事实性强、结构良好的语言。领域特化: 如果你正在构建一个旨在用于特定领域(例如医疗聊天机器人或代码生成助手)的模型,你可能会增加相关领域特定数据集(例如 PubMed 摘要或代码库)的权重。数据量与影响: 如果大型数据集(如 Common Crawl)按其原始大小比例抽样,它们很容易在训练组合中占据主导地位。分配明确的权重可以避免这种情况,确保较小、高质量或领域特定数据集仍能对模型的学习产生重要影响。防止数据污染: 如果在特定基准上进行评估,你可能会降低已知包含这些基准节选的数据来源的权重,或将其排除,以确保公平评估。期望的模型能力: 如果对话文本的流畅性很重要,你可能会包含并合理加权包含对话或论坛讨论的数据集。如果事实基础很重要,结构化文本和百科全书来源可能会获得更高权重。实现方面的考量实现来源加权通常涉及修改数据加载过程。数据加载器需要了解不同的来源及其相关权重,而不是从单一大型数据集中均匀抽样。在 PyTorch 中,你可以通过创建自定义的 IterableDataset 或使用能够适应多个底层数据集加权抽样的采样器来实现这一点。一个简化示例可能如下所示:import torch import numpy as np from torch.utils.data import IterableDataset, DataLoader # 假设这些是不同来源的占位数据集 # 实际中,这些会加载实际的 tokenized 数据 class DummyDataset(IterableDataset): def __init__(self, source_name, size): self.source_name = source_name self.size = size def __iter__(self): for i in range(self.size): # 生成虚拟数据:(样本数据, 来源标识符) yield torch.randn(512), self.source_name if i % 1000 == 0 and i > 0 : # 模拟可能的大型数据集 print(f"Yielded {i} from {self.source_name}") # 定义来源及其期望的抽样概率(权重) sources = { "web": {"dataset": DummyDataset("web", 1_000_000), "weight": 0.60}, "books": {"dataset": DummyDataset("books", 200_000), "weight": 0.15}, "code": {"dataset": DummyDataset("code", 300_000), "weight": 0.15}, "wiki": {"dataset": DummyDataset("wiki", 100_000), "weight": 0.10}, } source_names = list(sources.keys()) source_weights = np.array([sources[name]["weight"] for name in source_names]) # 归一化权重以确保它们总和为 1(如果已归一化则可选) # source_weights /= source_weights.sum() source_iters = {name: iter(sources[name]["dataset"]) for name in source_names} class WeightedSourceSampler(IterableDataset): def __init__(self, source_names, source_weights, source_iters): self.source_names = source_names self.source_weights = source_weights self.source_iters = source_iters def __iter__(self): while True: # 根据权重选择一个来源 chosen_source_name = np.random.choice( self.source_names, p=self.source_weights ) try: # 从所选来源的迭代器中获取下一个项 item = next(self.source_iters[chosen_source_name]) yield item except StopIteration: # 如果某个来源的迭代器已耗尽,重新创建它(或根据需要处理) print(f"正在重新启动 {chosen_source_name} 的迭代器") self.source_iters[chosen_source_name] = iter( sources[chosen_source_name]["dataset"] ) # 可选地,中断或为有限数据集实现逻辑 # 对于大规模预训练,迭代器通常会无限循环 # 重置后重新尝试获取: try: item = next(self.source_iters[chosen_source_name]) yield item except StopIteration: print( f"警告:{chosen_source_name} 的迭代器在重置后立即耗尽。" ) # 决定如何处理这种情况,例如重新抽样来源或引发错误 # 在此示例中,我们可能只是跳过此轮。 continue # 创建组合数据集采样器 weighted_sampler_dataset = WeightedSourceSampler( source_names, source_weights, source_iters ) # 与 DataLoader 一起使用 # 注意:num_workers > 0 需要仔细处理 # IterableDatasets 和迭代器 # 为简单起见,这里使用 num_workers=0。 # 实际实现需要多 worker 支持。 data_loader = DataLoader(weighted_sampler_dataset, batch_size=4, num_workers=0) # 获取一个批次示例 print("正在获取一个批次...") batch = next(iter(data_loader)) # batch[0] 包含数据张量,batch[1] 包含来源名称 print(f"批次数据形状:{batch[0].shape}") print(f"批次来源标识符:{batch[1]}") # 模拟获取更多批次以查看来源分布 print("\n正在获取更多批次...") for i in range(5): batch = next(iter(data_loader)) print(f"批次 {i+1} 来源:{batch[1]}")PyTorch 实现草图,用于根据预设权重从多个数据来源进行抽样。本示例演示了基本原理。大型语言模型的生产级数据加载器通常会涉及更复杂的机制,以提高效率、处理分布式训练以及管理无法完全载入内存的超大型数据集,可能还会使用流式传输技术。权衡与考量来源加权虽然功能强大,但也带来了复杂性:调优难度: 找到最佳权重通常需要大量试验,并在下游任务上进行评估,这可能会带来高昂的计算成本。偏差放大: 对某些来源过度加权可能导致模型采纳该数据中存在的偏见。反之,对重要但较小的数据集加权不足可能会阻碍模型在特定方面的表现。动态加权: 正如稍后在课程学习和退火的背景下所讨论的,权重可能不会在整个训练过程中固定不变。它们可能根据训练阶段或模型表现动态变化。来源加权是一种基本方法,用于管理大型语言模型预训练所用大规模数据集的构成。通过仔细考虑不同数据来源的质量、相关性和数据量,并分配适当的权重,工程师可以更好地指导学习过程,并塑造最终模型的能力。