Megatron-LM GPT2
如果您还没有,我们建议您在逐步完成本教程之前先阅读入门指南。
在本教程中,我们将向 Megatron-LM GPT2 模型添加 DeepSpeed,这是一个大型、强大的转换器。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 相比,我们将模型规模和速度提升到了一个全新的水平。
图 2:ZeRO-2 可扩展至 1700 亿个参数,吞吐量提高了 10 倍,实现了超线性加速,并通过避免对高达 130 亿个参数的模型进行代码重构来提高可用性。
更具体地说,DeepSpeed 和 ZeRO-2 在四个方面表现出色(如图 2 所示),支持大 10 倍的模型,速度提高 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 千兆浮点运算,聚合性能超过 15 拍浮点运算。对于相同大小的模型,与单独使用 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 亿个参数(用于比较)。
此外,在没有模型并行的情况下,这些模型可以在低带宽集群上进行训练,同时仍然可以实现比使用模型并行更好的吞吐量。例如,与在使用 40 Gbps Infiniband互连的四节点集群上使用模型并行相比,使用 ZeRO 支持的数据并行可以使 GPT-2 模型的训练速度提高近 4 倍,其中每个节点具有四个通过 PCI-E 连接的 16GB NVIDIA V100 GPU。因此,凭借这种性能提升,大型模型训练不再局限于具有超高速互连的 GPU 集群,也适用于带宽有限的中等规模集群。