Megatron-LM GPT2

如果您还没有，我们建议您在继续学习本教程之前先阅读入门指南。

在本教程中，我们将向 Megatron-LM GPT2 模型添加 DeepSpeed，这是一个大型且功能强大的 Transformer。Megatron-LM 支持模型并行和多节点训练。有关更多详细信息，请参阅相应的论文：Megatron-LM：使用模型并行训练数十亿参数的语言模型。

首先，我们讨论数据和环境设置以及如何使用原始 Megatron-LM 训练 GPT-2 模型。接下来，我们逐步启用该模型以使用 DeepSpeed 运行。最后，我们展示了使用 DeepSpeed 带来的**性能提升**和**内存占用减少**。

使用原始 Megatron-LM 训练 GPT-2

我们已将原始模型代码从Megatron-LM复制到 DeepSpeed Megatron-LM 并将其作为子模块提供。要下载，请执行

git submodule update --init --recursive

训练数据设置

  • 按照 Megatron 的说明下载 webtext 数据，并在 DeepSpeedExamples/Megatron-LM/data 下放置一个符号链接。

运行未修改的 Megatron-LM GPT2 模型

  • 对于单 GPU 运行
    • 更改 scripts/pretrain_gpt2.sh，将其 --train-data 参数设置为 "webtext"。
    • 运行 bash scripts/pretrain_gpt2.sh
  • 对于多个 GPU 和/或节点运行
    • 更改 scripts/pretrain_gpt2_model_parallel.sh
      • 将其 --train-data 参数设置为 "webtext"
      • GPUS_PER_NODE 表示测试中每个节点涉及的 GPU 数量
      • NNODES 表示测试中涉及的节点数量
    • 运行 bash scripts/pretrain_gpt2_model_parallel.sh

启用 DeepSpeed

要使用 DeepSpeed，我们将修改三个文件

  • arguments.py：参数配置
  • pretrain_gpt2.py：训练的主要入口点
  • utils.py：检查点保存和加载实用程序

参数解析

第一步是将 DeepSpeed 参数添加到 Megatron-LM GPT2 模型，方法是在 arguments.py 中使用 deepspeed.add_config_arguments()。

def get_args():
    """Parse all the args."""

    parser = argparse.ArgumentParser(description='PyTorch BERT Model')
    parser = add_model_config_args(parser)
    parser = add_fp16_config_args(parser)
    parser = add_training_args(parser)
    parser = add_evaluation_args(parser)
    parser = add_text_generate_args(parser)
    parser = add_data_args(parser)

    # Include DeepSpeed configuration arguments
    parser = deepspeed.add_config_arguments(parser)

初始化和训练

我们将修改 pretrain.py 以启用使用 DeepSpeed 进行训练。

初始化

我们使用 deepspeed.initialize 来创建 model_engine、optimizer 和 LR scheduler。以下是其定义

def initialize(args,
               model,
               optimizer=None,
               model_parameters=None,
               training_data=None,
               lr_scheduler=None,
               mpu=None,
               dist_init_required=True,
               collate_fn=None):

对于 Megatron-LM GPT2 模型，我们在其 setup_model_and_optimizer() 函数中初始化 DeepSpeed，如下所示，以传递原始 model、optimizer、args、lr_scheduler 和 mpu。

def setup_model_and_optimizer(args):
    """Setup model and optimizer."""

    model = get_model(args)
    optimizer = get_optimizer(model, args)
    lr_scheduler = get_learning_rate_scheduler(optimizer, args)

    if args.deepspeed:
        import deepspeed

        print_rank_0("DeepSpeed is enabled.")

        model, optimizer, _, lr_scheduler = deepspeed.initialize(
            model=model,
            optimizer=optimizer,
            args=args,
            lr_scheduler=lr_scheduler,
            mpu=mpu,
            dist_init_required=False
       )

请注意，当启用 FP16 时，Megatron-LM GPT2 会向 Adam 优化器添加一个包装器。DeepSpeed 有自己的 FP16 优化器，因此我们需要将 Adam 优化器直接传递给 DeepSpeed，而无需任何包装器。当启用 DeepSpeed 时，我们从 get_optimizer() 返回未包装的 Adam 优化器。

def get_optimizer(model, args):
    """Setup the optimizer."""

    ......

    # Use Adam.
    optimizer = Adam(param_groups,
                     lr=args.lr, weight_decay=args.weight_decay)

    if args.deepspeed:
        # fp16 wrapper is not required for DeepSpeed.
        return optimizer

使用训练 API

由 deepspeed.initialize 返回的 model 是我们将用于使用前向、后向和步骤 API 训练模型的DeepSpeed 模型引擎。

前向传播

前向传播 API 与 PyTorch 兼容，无需更改。

反向传播

反向传播是通过直接在模型引擎上调用 backward(loss) 来完成的。

    def backward_step(optimizer, model, lm_loss, args, timers):
        """Backward step."""

        # Total loss.
        loss = lm_loss

        # Backward pass.
        if args.deepspeed:
            model.backward(loss)
        else:
            optimizer.zero_grad()
            if args.fp16:
                optimizer.backward(loss, update_master_grads=False)
            else:
                loss.backward()

DeepSpeed 在使用小批量更新权重后会自动处理梯度归零。

此外，DeepSpeed 在内部处理了分布式数据并行和 FP16，简化了多个地方的代码。

(A) DeepSpeed 还会在梯度累积边界处自动执行梯度平均。因此，我们跳过 allreduce 通信。

        if args.deepspeed:
            # DeepSpeed backward propagation already addressed all reduce communication.
            # Reset the timer to avoid breaking timer logs below.
            timers('allreduce').reset()
        else:
            torch.distributed.all_reduce(reduced_losses.data)
            reduced_losses.data = reduced_losses.data / args.world_size
            if not USE_TORCH_DDP:
                timers('allreduce').start()
                model.allreduce_params(reduce_after=False,
                                       fp32_allreduce=args.fp32_allreduce)
                timers('allreduce').stop()

(B) 我们还跳过主梯度更新，因为 DeepSpeed 在内部处理了它。

        # Update master gradients.
        if not args.deepspeed:
            if args.fp16:
                optimizer.update_master_grads()

            # Clipping gradients helps prevent the exploding gradient.
            if args.clip_grad > 0:
                if not args.fp16:
                    mpu.clip_grad_norm(model.parameters(), args.clip_grad)
                else:
                    optimizer.clip_master_grads(args.clip_grad)

        return lm_loss_reduced
更新模型参数

DeepSpeed 引擎中的 step() 函数会更新模型参数以及学习率。

     if args.deepspeed:
         model.step()
     else:
         optimizer.step()

         # Update learning rate.
         if not (args.fp16 and optimizer.overflow):
             lr_scheduler.step()
         else:
             skipped_iter = 1
损失缩放

GPT2 训练脚本在训练期间记录损失缩放值。在 DeepSpeed 优化器中，此值存储为 cur_scale 而不是 Megatron 优化器中的 loss_scale。因此，我们在日志记录字符串中适当地替换了它。

             if args.fp16:
                 log_string += ' loss scale {:.1f} |'.format(
                     optimizer.cur_scale if args.deepspeed else optimizer.loss_scale)

检查点保存和加载

DeepSpeed 引擎具有灵活的 API 用于检查点保存和加载，以处理来自客户端模型及其自身内部的状态。

def save_checkpoint(self, save_dir, tag, client_state={})
def load_checkpoint(self, load_dir, tag)

要使用 DeepSpeed，我们需要更新 utils.py，Megatron-LM GPT2 在其中保存和加载检查点。

创建一个名为 save_ds_checkpoint() 的新函数，如下所示。新函数会收集客户端模型状态并通过调用 DeepSpeed 的 save_checkpoint() 将其传递给 DeepSpeed 引擎。

 def save_ds_checkpoint(iteration, model, args):
     """Save a model checkpoint."""

     sd = {}
     sd['iteration'] = iteration
     # rng states.
     if not args.no_save_rng:
         sd['random_rng_state'] = random.getstate()
         sd['np_rng_state'] = np.random.get_state()
         sd['torch_rng_state'] = torch.get_rng_state()
         sd['cuda_rng_state'] = get_accelerator().get_rng_state()
         sd['rng_tracker_states'] = mpu.get_cuda_rng_tracker().get_states()

     model.save_checkpoint(args.save, iteration, client_state = sd)

在 Megatron-LM GPT2 的 save_checkpoint() 函数中，添加以下几行以调用上述函数以使用 DeepSpeed。

 def save_checkpoint(iteration, model, optimizer,
                     lr_scheduler, args):
     """Save a model checkpoint."""
     if args.deepspeed:
         save_ds_checkpoint(iteration, model, args)
     else:
		......

在 load_checkpoint() 函数中，使用 DeepSpeed 检查点加载 API，如下所示，并返回客户端模型的状态。

 def load_checkpoint(model, optimizer, lr_scheduler, args):
     """Load a model checkpoint."""

     iteration, release = get_checkpoint_iteration(args)

     if args.deepspeed:
         checkpoint_name, sd = model.load_checkpoint(args.load, iteration)

         if checkpoint_name is None:
             if mpu.get_data_parallel_rank() == 0:
                 print("Unable to load checkpoint.")
             return iteration
     else:
         ......

DeepSpeed 激活检查点（可选）

DeepSpeed 可以通过将激活检查点跨模型并行 GPU 分区或将其卸载到 CPU 来减少模型并行训练期间的激活内存。这些优化是可选的，可以跳过，除非激活内存成为瓶颈。要启用分区激活，我们使用 deepspeed.checkpointing API 替换 Megatron 的激活检查点和随机状态跟踪器 API。替换应在首次调用这些 API 之前进行。

a) 在 pretrain_gpt.py 中替换

    # Optional DeepSpeed Activation Checkpointing Features
    #
    if args.deepspeed and args.deepspeed_activation_checkpointing:
        set_deepspeed_activation_checkpointing(args)

def set_deepspeed_activation_checkpointing(args):

    deepspeed.checkpointing.configure(mpu,
                            deepspeed_config=args.deepspeed_config,
                            partition_activation=True)

    mpu.checkpoint = deepspeed.checkpointing.checkpoint
    mpu.get_cuda_rng_tracker = deepspeed.checkpointing.get_cuda_rng_tracker
    mpu.model_parallel_cuda_manual_seed =
                    deepspeed.checkpointing.model_parallel_cuda_manual_seed

b) 在 mpu/transformer.py 中替换

if deepspeed.checkpointing.is_configured():
    global get_cuda_rng_tracker, checkpoint
    get_cuda_rng_tracker = deepspeed.checkpoint.get_cuda_rng_tracker
    checkpoint = deepspeed.checkpointing.checkpoint

通过这些替换，可以使用 deepspeed.checkpointing.configure 或在 deepspeed_config 文件中指定各种 DeepSpeed 激活检查点优化，例如激活分区、连续检查点和 CPU 检查点。

训练脚本

我们假设 webtext 数据已在上一步中准备完毕。要开始使用 DeepSpeed 训练 Megatron-LM GPT2 模型，请执行以下命令以开始训练。

  • 单 GPU 运行
    • 运行 bash scripts/ds_pretrain_gpt2.sh
  • 多个 GPU/节点运行
    • 运行 bash scripts/ds_zero2_pretrain_gpt2_model_parallel.sh

使用 GPT-2 进行 DeepSpeed 评估

DeepSpeed 通过高级ZeRO 优化器有效地启用大型模型的训练。2020 年 2 月，我们在 DeepSpeed 中发布了 ZeRO 优化器的一个子集，它执行优化器状态分区。我们将其称为 ZeRO-1。2020 年 5 月，我们在 DeepSpeed 中扩展了 ZeRO-1，其中包括 ZeRO 的其他优化，包括梯度和激活分区，以及连续内存优化。我们将其称为 ZeRO-2。

ZeRO-2 显著降低了训练大型模型的内存占用，这意味着可以使用以下方式训练大型模型：i) 较少的模型并行性；ii) 较大的批量大小。较低的模型并行性程度通过提高计算粒度（例如矩阵乘法，其中性能直接与矩阵大小相关）来提高训练效率。此外，较少的模型并行性还会导致模型并行 GPU 之间的通信减少，从而进一步提高性能。较大的批量大小具有类似的效果，即提高计算粒度并减少通信，从而也带来更好的性能。因此，通过将 DeepSpeed 和 ZeRO-2 集成到 Megatron 中，与仅使用 Megatron 相比，我们将模型规模和速度提升到一个全新的水平。

DeepSpeed-vs-Megatron

图 2：ZeRO-2 可扩展到 1700 亿个参数，吞吐量高出 10 倍，获得超线性加速，并通过避免对高达 130 亿个参数的模型进行代码重构来提高可用性。

更具体地说，DeepSpeed 和 ZeRO-2 在四个方面（如图 2 所示）表现出色，支持大一个数量级的模型，速度快 10 倍，具有超线性可扩展性，以及改进的可用性，从而使大型模型训练民主化。这四个方面将在下面详细介绍。

**模型大小**：OpenAI GPT-2、NVIDIA Megatron-LM、Google T5 和 Microsoft Turing-NLG 等最先进的大型模型的参数大小分别为 15 亿、83 亿、110 亿和 170 亿。ZeRO-2 提供系统支持以有效地运行 1700 亿个参数的模型，比这些最大的模型大一个数量级（图 2，左上角）。

**速度**：改进的内存效率带来更高的吞吐量和更快的训练速度。图 2（左下角）显示了 ZeRO-2 和 ZeRO-1（两者都将 ZeRO 驱动的并行数据与 NVIDIA Megatron-LM 模型并行相结合）的系统吞吐量，以及使用最先进的模型并行方法 Megatron-LM 本身（图 2，左下角的基线）。ZeRO-2 在具有 400 个 NVIDIA V100 GPU 的集群上运行 1000 亿参数模型，每个 GPU 的吞吐量超过 38 teraflops，总性能超过 15 petaflops。对于相同大小的模型，与仅使用 Megatron-LM 相比，ZeRO-2 的训练速度快 10 倍；与 ZeRO-1 相比，训练速度快 5 倍。

**可扩展性**：我们观察到超线性加速（图 2，右上角），其中当 GPU 数量加倍时，性能超过两倍。当我们提高并行数据程度时，ZeRO-2 会减少模型状态的内存占用，从而允许我们为每个 GPU 适应更大的批量大小，从而带来更好的性能。

大模型训练的民主化：ZeRO-2 使模型科学家能够高效地训练高达 130 亿参数的模型，而无需通常需要模型重构的模型并行（图 2，右下角）。130 亿参数大于大多数最先进的大模型（例如 Google T5，参数为 110 亿）。因此，模型科学家可以自由地试验大型模型，而无需担心模型并行。相比之下，经典数据并行方法（如 PyTorch 分布式数据并行）的实现使用 14 亿参数模型时会耗尽内存，而 ZeRO-1 则支持最多 60 亿参数。

此外，在没有模型并行的情况下，这些模型可以在低带宽集群上进行训练，同时仍然比使用模型并行获得显着更高的吞吐量。例如，GPT-2 模型在使用 ZeRO 支持的数据并行时，与在使用四个具有 40 Gbps InfiniBand 相互连接的节点的集群上使用模型并行相比，训练速度快近 4 倍，其中每个节点具有四个通过 PCI-E 连接的 NVIDIA 16GB V100 GPU。因此，凭借这种性能提升，大型模型训练不再局限于具有超高速互连的 GPU 集群，也适用于带宽有限的中等规模集群。

更新：