课程学习：一种用于高效稳定地预训练十亿级GPT模型的正则化方法

注意！ 在 2022 年 12 月 12 日，我们发布了 DeepSpeed 数据效率库，它提供了更通用的课程学习支持。下面的这个旧版课程学习功能仍然受支持，但我们建议使用数据效率库（教程）。

注意： 本教程于 2021 年 10 月 29 日更新。更改包括：1) 更详细的调优策略。2) 支持流水线并行。3) 基于 Token 的学习率衰减。4) github.com/deepspeedai/Megatron-DeepSpeed 上一个新的 GPT-2 示例。详见下文。

在本教程中，我们将介绍 DeepSpeed 基于课程学习的数据管道，它在训练早期呈现更容易或更简单的示例。通过实现 8 倍/4 倍更大批大小/学习率的稳定训练（而基线方法则在训练发散时遇到困难），我们观察到课程学习（基于序列长度）提供了稳定且 3.3 倍更快的 GPT-2 预训练（在 1.17 亿和 15 亿参数上测试），同时具有更好的基于 Token 的收敛速度和零样本 WikiText-103/LAMBADA 评估结果。此外，由于课程学习只影响数据管道，其优势与许多 DeepSpeed 功能和其他系统优化技术互补。例如，课程学习兼容 DeepSpeed 的 ZeRO 冗余优化器、ZeRO-Offload 和 3D 并行。

为了说明课程学习的优点和用法，我们以 Megatron-LM GPT-2 预训练任务为例。有关此任务的更多详细信息，请参阅 Megatron-LM GPT2 教程。此外，我们还有一篇论文提供了包括实现和评估在内的技术细节。

1. 配置和调优策略

通过在 DeepSpeed 配置文件中设置 curriculum_learning 键可以使用课程学习

{
  "train_batch_size": 4096,
  "gradient_accumulation_steps": 1,
  "steps_per_print": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.00015,
      "max_grad_norm": 1.0,
      "betas": [0.9, 0.95]
    }
  },
  "gradient_clipping": 1.0,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "consecutive_hysteresis": false,
    "min_loss_scale": 1
  },
  "curriculum_learning": {
    "enabled": true,
    "curriculum_type": "seqlen",
    "min_difficulty": 8,
    "max_difficulty": 1024,
    "schedule_type": "fixed_linear",
    "schedule_config": {
      "total_curriculum_step": 15000,
      "difficulty_step": 8
    }
  }
}

为了支持课程学习，我们添加了以下新参数

curriculum_type 是课程难度指标的类型。目前我们支持 seqlen 指标，它在训练早期呈现较短的序列。我们通过在实际前向传播之前执行训练数据序列截断来实现这种类型的课程学习。我们将在下面的 Megatron-LM GPT-2 预训练示例中描述如何实现这一点。

min_difficulty 是起始难度级别。对于 seqlen 指标，这意味着我们从序列长度为 min_difficulty 开始。我们观察到较低的 min_difficulty 通常能提供更好的稳定性/收敛速度，但有两个注意事项：首先，有时（特别是对于大型模型）从太小的难度级别开始可能会导致严重的过拟合（例如，训练损失发散或验证困惑度波动），从而损害收敛。其次，对于 seqlen 指标，我们建议将 min_difficulty 设置为 8（对于 FP16 数据）或 16（对于 INT8 数据）的倍数，以启用 NVIDIA GPU 的 Tensor Core 加速。为了调优 seqlen 指标的这个超参数，我们建议从 min_difficulty 为 8（百万级模型）或 64（十亿级模型）开始，如果一开始就观察到发散或验证困惑度波动，则增加它。

max_difficulty 是结束难度级别。对于 seqlen 指标，应将其设置为完整序列长度（例如，Megatron-LM GPT-2 预训练的 1024）。

schedule_type 是课程学习的调度策略（即在特定步骤使用哪个难度级别）。目前我们支持三种调度：fixed_linear、fixed_root 和 fixed_discrete。我们建议首先尝试 fixed_linear 调度，它更容易调优，并在我们的测试中提供了极大的训练稳定性/效率增益。每种调度都有其自己的配置

1.1 fixed_linear 调度

对于 fixed_linear 调度，有两个配置

"schedule_type": "fixed_linear",
"schedule_config": {
  "total_curriculum_step": 15000,
  "difficulty_step": 8
}

total_curriculum_step 是课程学习的总步数。对于 fixed_linear 调度，难度级别将在 total_curriculum_step 步内从 min_difficulty 线性增加到 max_difficulty。此配置必须针对每个训练任务进行调优。我们观察到 total_curriculum_step 过小和过大都是次优的：total_curriculum_step 过小可能导致课程学习无法提供足够的训练稳定性益处，从而训练仍可能发散；total_curriculum_step 过大可能导致模型在课程学习期间对更容易/更简单的训练数据过拟合，从而损害整体收敛。为了调优此超参数，我们建议进行二分查找以找到在学习率预热前几倍步数内没有显著验证困惑度波动的最大 total_curriculum_step。其基本原理可在我们的论文附录 A.1 中找到。

difficulty_step 配置确保任何时候的难度级别都是 difficulty_step 的倍数。较小的值更佳，因为它提供更平滑的课程和更好的稳定性。我们通常将其设置为 8（对于 FP16 数据）或 16（对于 INT8 数据），以启用 NVIDIA GPU 的 Tensor Core 加速。如果这与您的硬件无关，您可以将其设置为 1。

1.2 fixed_root 调度

对于 fixed_root 调度，有三个配置

"schedule_type": "fixed_root",
"schedule_config": {
  "total_curriculum_step": 15000,
  "difficulty_step": 8,
  "root_degree": 2
}

total_curriculum_step 和 difficulty_step 的含义与 fixed_linear 调度相同。 root_degree 决定了调度根函数的根次方。特定步骤的难度级别确定为 ((current step/total_curriculum_step)**(1/root_degree)) * (max_difficulty - min_difficulty) + min_difficulty。因此，fixed_linear 基本上是 fixed_root 的一个特例，其中 root_degree 为 1。在我们（有限的）研究中，我们发现 fixed_root 调度相对于 fixed_linear 调度没有提供任何明显的优势，同时还需要一个额外的参数。

1.3 fixed_discrete 调度

对于 fixed_discrete 调度，有两个配置

"schedule_type": "fixed_discrete",
"schedule_config": {
  "difficulty": [1,2,3],
  "max_step": [5,10]
}

difficulty 是一个难度级别列表，用于调度期间。 max_step 是一个步数时间戳列表，用于确定何时切换到下一个难度级别。例如，上面的 json 配置意味着在步骤 1-5 使用难度 1，在步骤 6-10 使用难度 2，从步骤 11 开始使用难度 3。这种 fixed_discrete 调度提供了最灵活的课程学习调度。然而，我们发现这种调度的一个风险是，如果模型在某个难度级别停留时间过长，切换到下一个难度时可能会因为严重的过拟合而发生训练发散。

2. Megatron-LM GPT-2 预训练的课程学习

注意！ 在 2021 年 10 月 29 日更新后，现在 Megatron-LM GPT-2 预训练有两个课程学习示例。它们都具有一些独特的功能和限制。详见下文。

我们提供了两个 Megatron-LM GPT-2 预训练的课程学习示例

第一个位于 Megatron-DeepSpeed/tree/main/examples_deepspeed/curriculum_learning。此集成基于较新的 Megatron-LM 分支，并且只有此课程学习示例支持流水线并行。然而，截至 2021 年 10 月 29 日，我们尚未在此分支上验证 ZeRO-2 和 ZeRO-3。总的来说，如果您的模型不需要 ZeRO-2/3，我们强烈建议您使用此示例。

第二个位于 DeepSpeedExamples/Megatron-LM-v1.1.5-ZeRO3/curriculum_learning/。此集成基于一个我们最终将弃用的旧版 Megatron-LM 硬拷贝，并且此课程学习示例不支持流水线并行。我们建议您仅在模型需要 ZeRO-2/3 时使用此示例。

除了上面描述的 DeepSpeed 课程学习 JSON 配置之外，用户端还需要进行一些其他必要的更改才能集成课程学习

2.1 训练数据截断

为了启用基于 seqlen 的课程学习，我们需要根据给定的课程序列长度添加训练数据截断功能。对于不使用流水线并行的情况，有必要在模型的前向传播中添加 curriculum_seqlen 参数并使用它来执行训练数据序列长度截断。对于 Megatron-LM GPT-2 预训练，我们在 megatron/model/gpt2_model.py 的 forward() 和 pretrain_gpt2.py 的 forward_step() 中实现了这一点。

对于使用流水线并行的情况，由于 DeepSpeed 引擎的限制，我们无法在前向传播中注入 curriculum_seqlen 参数。相反，我们在用户端创建了 deepspeed.runtime.data_pipeline.curriculum_scheduler 的一个副本，并使用它来检索 curriculum_seqlen。此实现在 megatron/training.py 中可以找到。

2.2 禁用批大小预热 (--rampup-batch-size)

在我们的论文第 5.4 节中，我们证明了课程学习（基于 seqlen）提供了比 Open AI GPT-3 引入的批大小预热技术更好的训练稳定性。因此，在使用课程学习时，您需要从训练脚本中删除 --rampup-batch-size 配置。不建议同时使用课程学习和批大小预热，因为它们都会减少批次中的 Token 数量。您可能想要进行的另一个相关更改是增加您的微批大小，因为没有批大小预热，您的批大小现在将是固定的。

2.3 基于 Token 的训练终止

因为课程学习在训练过程中改变了每个序列/样本的长度，所以很难/不可能使用步数/样本数来精确地在所需的 Token 数量处终止训练。因此，我们添加了一个 --train-tokens 配置用于精确的基于 Token 的终止。我们建议将您原来的 --train-samples 或 --train-iters 增加到足够大的数量（例如，基线使用的 3 倍），并将 --train-tokens 设置为精确所需的训练 Token 数量。

2.4 基于 Token 的学习率衰减

再次，因为课程学习改变了每个批次中 Token 的数量，在我们的论文附录 A.2 中，我们展示了将学习率衰减改为基于 Token 的也是必要的（以避免学习率衰减过快）。因此，我们添加了一个 --lr-decay-tokens，它将是学习率衰减 Token 的数量。如果之前您使用 --lr-decay-samples，您可以通过将前者乘以完整 seqlen（例如，GPT-2 为 1K，GPT-3 为 2K）来计算您的 --lr-decay-tokens。如果之前您使用 --lr-decay-iters，您可以通过将前者乘以完整 seqlen 和全局批大小来计算您的 --lr-decay-tokens。然后，您需要在脚本中将 --lr-decay-samples 或 --lr-decay-iters 替换为 --lr-decay-tokens。

2.5 学习率预热调整

对于学习率预热，我们不将其更改为基于 Token 的，因为这样做对于课程学习意味着减慢学习率预热，这既不必要也有害。然而，为了避免过快的预热，您可能需要根据非 CL 情况调整您的 --lr-warmup-samples 或 --lr-warmup-iters（例如，如果您在非 CL 情况下使用了 --rampup-batch-size，对于 CL 我们不使用它，因此每个批次的样本数量在开始时会有所不同）。假设您想使用 X 个 Token 进行学习率预热（对于 OpenAI GPT-3，这是 3.75 亿个 Token），那么对于课程学习情况，您应该将 --lr-warmup-samples 设置为 X 除以 min_difficulty，或者将 --lr-warmup-iters 设置为 X 除以 min_difficulty * --global-batch-size。这是一个粗略的估计，基于课程学习从 seqlen min_difficulty 开始，并且在学习率预热期间不会增加太多。

更新时间：