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

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

注意： 本教程于 2021 年 10 月 29 日更新。更改包括：1）更详细的调整策略。2）流水线并行支持。3）基于 Token 的学习率衰减。4）github.com/microsoft/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 配置。不建议同时使用课程学习和批大小预热，因为它们都会减少批次中的标记数量。您可能需要进行的另一个相关更改是增加微批次大小，因为在没有批大小预热的情况下，您的批次大小现在将是固定的。

2.3 基于 Token 的训练终止

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

2.4 基于 Token 的学习率衰减

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

2.5 学习率预热调整

对于 LR 预热，我们不会将其更改为基于标记的，因为对于课程学习这样做意味着减慢 LR 预热速度，这既没有必要，也弊大于利。但是，为了避免预热过快，您可能需要根据各种原因调整非 CL 情况下的 --lr-warmup-samples 或 --lr-warmup-iters（例如，如果您在非 CL 情况下使用了 --rampup-batch-size，则对于 CL 我们不使用它，因此每个批次的样本数量在开始时将不同）。假设您希望使用 X 个标记来预热 LR（对于 OpenAI GPT-3，这是 375M 个标记），那么对于课程学习案例，您应该将 --lr-warmup-samples 设置为 X 除以 min_difficulty，或将 --lr-warmup-iters 设置为 X 除以 min_difficulty * --global-batch-size。这是一个粗略的估计，基于课程学习从 seqlen min_difficulty 开始，并且在 LR 预热期间不会增加太多。