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

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

注意： 本教程已于2021年10月29日更新。更新内容包括：1）更详细的调优策略。2）管道并行支持。3）基于令牌的学习率衰减。4）在github.com/microsoft/Megatron-DeepSpeed上新增了一个GPT-2示例。详情见下文。

在本教程中，我们介绍了DeepSpeed基于课程学习的数据管道，该管道在训练早期提供更容易或更简单的示例。通过启用8倍/4倍更大的批量大小/学习率进行稳定训练（而基线方法在训练发散方面存在困难），我们观察到基于序列长度的课程学习提供了稳定且3.3倍更快的GPT-2预训练（在117M和1.5B参数上测试），同时具有更好的逐词收敛速度和零样本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计划，有两种配置：

"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计划，有三种配置：

"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计划，有两种配置：

"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 基于令牌的训练终止

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

2.4 基于Token的LR衰减

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

2.5 LR预热调整

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

更新: