doremi

doremi？：使用最小值优化的域重新加权

Doremi的Pytorch实现，一种用于优化语言建模数据集的数据混合物的算法。现代大型语言模型在许多领域（网络，书籍，ARXIV等）上进行了培训，但是每个要训练的域都不清楚，尤其是因为这些模型将用于各种下游任务（没有特定的目标分配以优化）。 Doremi使用分布强大的优化（DRO）调整数据混合物以鲁棒至目标分布。 Doremi使用DRO训练一个小的代理模型，该模型与预验证的参考模型相比，该模型基于代理模型的多余损失，动态重量或减肥域。参考模型提供了可实现的最佳损失的估计，以避免对高熵 /硬域的悲观情绪。然后，调谐的数据混合物可用于更有效地训练更大的模型。在本文中，280m代理模型可以改善8B参数模型的训练（大30倍），从而使其能够更快地达到基线8B性能2.6倍。下图提供了Doremi的概述。查看纸张以获取更多详细信息。

作为黑匣子，该代码库给定文本数据集输出优化的域权重。其他一些有用的组件：具有域级加权采样，简单的下游评估线束和HuggingFace Trainer + FlashAttention2集成的快速，可重新恢复的数据加载器。

入门

要开始，请克隆回购并安装它：

 git clone [email protected]:/sangmichaelxie/doremi.git
pip install -e doremi
cd doremi && bash scripts/setup_flash.sh

scripts/setup_flash.sh中的汇编可能需要大量时间（小时）。所有代码都应从最外层的doremi目录运行。在开始之前，请在此存储库的外部目录中的constants.sh文件中写入高速缓存目录，数据目录等。您还可以在此处放置任何CONDA或VIRTAILENV激活命令。这是constants.sh文件内容的示例（以称为sample_constants.sh的文件提供）：

 #!/bin/bash
CACHE=/path/to/cache
DOREMI_DIR=/path/to/this/repo
PILE_DIR=/path/to/pile
PREPROCESSED_PILE_DIR=/path/to/preprocessed  # will be created by scripts/run_preprocess_pile.sh
MODEL_OUTPUT_DIR=/path/to/model_output_dir
WANDB_API_KEY=key  # Weights and Biases key for logging
PARTITION=partition # for slurm
mkdir -p ${CACHE}
mkdir -p ${MODEL_OUTPUT_DIR}
source ${DOREMI_DIR}/venv/bin/activate  # if you installed doremi in venv

这是如何运行示例脚本以在堆上进行数据预处理的示例脚本，该脚本将桩的数据分隔为域并将其分开：

 bash scripts/run_preprocess_pile.sh

这是一个示例脚本，可以运行1200m基线，代理和主型号（Doremi Pipeline中的所有3个步骤），并在一个带有8个A100 GPU的一个节点上测试。这是论文中桩实验的小版本。该脚本将自动运行困惑和几次评估：

 bash scripts/run_pile.sh

这些脚本按照论文运行200k步骤。 doremi运行使用文件名<RUN_NAME>.json中的configs目录中的域权重。

在您自己的数据集上运行doremi

要在您自己的数据集上运行doremi，请以以下格式提供预处理（令牌化）数据：

 top_level/
    domain_name_1/
        files...
    domain_name_2/
        files...
    ...

每个内部目录（例如， domain_name_1 ）可以通过huggingface的load_from_disk方法加载。如果您的数据以不同的格式，则可以在doremi/dataloader.py中添加自定义数据加载功能。您还需要编写一个配置文件并将其保存到configs/和写入运行脚本类似于scripts/runs/run_pile_baseline120M.sh and scripts/runs/run_pile_doremi120M.sh请参考配置文件。配置文件指定从域名到混合物重量的映射。名称不必按顺序（Doremi始终先对域名进行排序以确定固定订购），并且不必将权重标准化。

提示和详细信息

• 参考域权重的选择：通常，参考域权重是表达对域重要性的先验的一种方式。一个合理的默认值是根据每个域的大小设置参考模型的域权重，这确保了小域不会过分代表（并且过于拟合）。如果某些域特别重要，则当然可以增加其相应的参考域重量。均匀的域重量可以用作参考域的权重，以避免将任何先前放在域的权重上，但是可能需要迭代的doremi（见下文）。
• 迭代的Doremi ：在某些情况下，您可能需要运行超过1轮的doremi。为了运行迭代的doremi，我们使用上一轮优化的域权重训练新的参考模型。尤其是如果您从次优的参考域重量（例如均匀的权重）开始，使用这些权重训练的参考模型也将是次优的。
• 域重量的更新速率（ --reweight_eps ） ：纸张中使用的默认设置为1，尽管可以针对不同的数据集进行调整。通常，我们预计训练期间的域重量会有些嘈杂，平均域的重量大多是光滑的。
• 代币仪的选择：我们发现，较新的令牌剂（如Neox）倾向于在较小的尺度上提供更好的模型性能，因此我们建议使用这些较旧的令牌（例如GPT2）。
• 在域的子集上运行：默认数据集加载代码只会加载配置文件中列出的域。可以通过从配置文件中删除域来实现在域的子集上。
• 梯度积累：梯度积累应起作用，在初步测试中，该行为类似于无蓄能情况。但是，目前在此实施中使用积累并不等于没有积累。如果我们积累了k步的梯度，则将针对先前迭代的陈旧域重量计算出k-1梯度（ k=1不存在此问题）。
• 多节点培训：我们目前不支持多节点培训，而是专注于单节点，多GPU设置。

请注意，该回购和论文之间存在一些差异，该论文是在Google上开发的，即：

• Pytorch vs jax
• 模型体系结构的细微差异
• 使用的象征器（纸中使用的256K词汇尺寸，而标准的开源标记器大约为50k词汇尺寸）。这可以显着影响由令牌计数计算的数据混合物。您应该在自己的特定培训设置中运行Doremi，以获得最佳结果。

样本运行结果

以下是使用120m代理和参考模型（带有scripts/run_pile.sh ）在桩上的一轮Doremi的结果。我们使用优化的权重训练120m型号（ configs/pile_doremi_r1_120M_ref:pile_baseline_50kvocab_nopack_120M.json ），并将其与基线（灰色）进行比较。两个基线代表了两种略有不同的方法来计算桩中的基线域权重（Nopack计算每个文档将每个文档填充到上下文窗口长度后的每个域中的示例数，而PACK首先将文档串联到一个域中），这会产生相似的模型。经过doremi域重量训练的模型在训练期间很早就超过了基线的单杆性能，在所有任务中，在70k步长（快3倍）内。 DOREMI模型超过了20K步骤内的平均基线单球性能，在15/22域上改善或可比性的困惑，并改善了整个域的平均平均和最差的困惑。

引用

如果这对您有用，请引用纸张：

 @article{xie2023doremi,
  author = {Sang Michael Xie and Hieu Pham and Xuanyi Dong and Nan Du and Hanxiao Liu and Yifeng Lu and Percy Liang and Quoc V. Le and Tengyu Ma and Adams Wei Yu},
  journal = {arXiv preprint arXiv:2305.10429},
  title = {DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining},
  year = {2023},
}