尧图网站建设 
在 Megatron-LM GPT-2 上接入 DeepSpeed模型并行训练的逐文件改造与 ZeRO-2 实践【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文是一份面向工程实战的技术指南完整讲解如何把 DeepSpeed 分布式训练优化库接入基于 Megatron-LM 的 GPT-2 大模型训练代码从数据与环境准备、原始模型跑通到逐文件改造参数解析、引擎初始化、前向/反向/更新循环与断点续训再到可选激活检查点优化与 ZeRO-2 规模化训练效果。读完本文你将能够复现官方教程中的完整接入流程并理解每一步背后 DeepSpeed 引擎Engine的源码级实现原理。本文内容以仓库内文档 docs/_tutorials/megatron.md 为骨架并结合本仓库中 DeepSpeed 核心源码展开印证。背景与适用范围什么样的模型适合走这条接入路径Megatron-LM 是一套支持**模型并行model parallelism与多节点multi-node**训练的 Transformer 训练框架其对应论文为Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism。GPT-2 是其中典型的超大规模 Transformer 模型。DeepSpeed 将 ZeRO 数据并行与 Megatron 的模型并行相结合后可以解决「模型大到单卡乃至单节点放不下」的训练难题。本教程面向的正是此类已具备模型并行能力、希望在不大规模改写逻辑的前提下叠加 ZeRO 优化的训练代码。需要先说明一个代码归属问题避免混淆本教程提到的megatron/arguments.py、pretrain_gpt.py、megatron/utils.py、mpu/transformer.py以及各类.sh启动脚本都位于Megatron-LM 训练代码仓库DeepSpeed 社区维护的 Megatron-DeepSpeed 分叉作为示例仓库的 submodule 引入而不是DeepSpeed 核心库所在目录中的文件。DeepSpeed 核心库即本仓库deepspeed/目录为接入过程提供的是deepspeed.initialize、deepspeed.add_config_arguments、Engine 的backward/step/save_checkpoint/load_checkpoint与激活检查点等 API。因此本仓库内并找不到megatron/目录需要先按下文拉取 Megatron 代码。如果你还没有阅读过入门指南建议先浏览 Getting Started 教程其中覆盖了 DeepSpeed 的安装、基础配置与最小训练示例。第一步用原始 Megatron-LM 跑通 GPT-2改造基线在接入 DeepSpeed 之前先在原始代码上跑通训练作为后续性能对比的基线。拉取 Megatron-LM 训练代码Megatron-LM 的原版模型代码已被复制到 DeepSpeed 的 Megatron 分叉中并以 submodule 形式提供。执行以下命令获取git submodule update --init --recursive训练数据准备下载 OpenAIwebtext数据集按 Megatron 官方 README 中Collecting GPT2 WebText Data一节提供的步骤下载数据并在DeepSpeedExamples/Megatron-LM/data目录下为数据放置一个符号链接使训练脚本可以按data路径找到语料。运行未修改的 Megatron-LM GPT-2单 GPU 运行修改examples/pretrain_gpt.sh把它的--train-data参数设为webtext然后执行bash examples/pretrain_gpt.sh多 GPU / 多节点运行修改examples/pretrain_gpt_distributed_with_mp.sh将--train-data参数设为webtextGPUS_PER_NODE表示测试中每个节点投入的 GPU 数量NNODES表示参与测试的节点数然后执行bash examples/pretrain_gpt_distributed_with_mp.sh启用 DeepSpeed 的总体改造清单使用 DeepSpeed 需要修改 Megatron-LM 训练代码中的三个文件分工如下文件职责megatron/arguments.py参数配置把 DeepSpeed 命令行参数注册进参数解析器pretrain_gpt.py训练主入口初始化 DeepSpeed 引擎并改造训练循环megatron/utils.py模型检查点的保存与加载逻辑对接 DeepSpeed修改一把 DeepSpeed 参数注入参数解析器DeepSpeed 提供了一个便捷函数deepspeed.add_config_arguments(parser)用于向已有的argparse.ArgumentParser注入一套 DeepSpeed 命令行参数。在megatron/arguments.py的get_args()末尾调用它def get_args(): Parse all the args. parser argparse.ArgumentParser(descriptionPyTorch 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)从源码看add_config_arguments()实际注入的核心参数包括--deepspeed布尔开关表示启用 DeepSpeed--deepspeed_config json 文件路径DeepSpeed 运行时 JSON 配置文件路径。正是因为这一处注册后续训练脚本里才可以通过--deepspeed与--deepspeed_config显式打开 DeepSpeed 并指定配置。修改二初始化 DeepSpeed 引擎并改造训练循环deepspeed.initialize是 DeepSpeed 引擎的入口它同时负责创建model engine、optimizer 与 LR scheduler。结合当前仓库 deepspeed/init.py 中的实际签名可见def initialize(args, model, optimizerNone, model_parametersNone, training_dataNone, lr_schedulerNone, mpuNone, dist_init_requiredTrue, collate_fnNone):model任何 wrapper 之前的原始nn.Moduleoptimizer用户自定义 optimizer若提供则覆盖 JSON 配置中的 optimizer 定义mpu提供get_{model,data}_parallel_{rank,group,world_size}()接口的模型并行单元对象用于把 Megatron 的模型并行信息告知 DeepSpeeddist_init_required是否由 DeepSpeed 自动初始化 torch 分布式环境。对于 Megatron-LM GPT-2在其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( modelmodel, optimizeroptimizer, argsargs, lr_schedulerlr_scheduler, mpumpu, dist_init_requiredFalse )注意这里dist_init_requiredFalseMegatron 训练代码本身已经完成了分布式环境的初始化含mpu建立的各种进程组因此需要明确告知 DeepSpeed 不要再重复初始化。一个重要细节FP16 下必须传「未包装」的 AdamMegatron-LM 在启用 FP16 时会给Adam优化器套一层自己的 FP16 wrapper负责维护 master weights / master grads 与动态 loss scaling。但DeepSpeed 拥有自己的 FP16 Optimizer因此当args.deepspeed为真时get_optimizer()必须返回未包装的裸Adam把 FP16 相关管理全部交给 DeepSpeeddef get_optimizer(model, args): Setup the optimizer. ...... # Use Adam. optimizer Adam(param_groups, lrargs.lr, weight_decayargs.weight_decay) if args.deepspeed: # fp16 wrapper is not required for DeepSpeed. return optimizer使用 Training API前向 / 反向 / 参数更新deepspeed.initialize返回的model是DeepSpeed Model EngineDeepSpeedEngine训练将围绕 engine 暴露的 forward、backward、step API 进行。前向传播前向传播 API 与 PyTorch 完全兼容不需要任何改动。反向传播反向传播直接调用 engine 的model.backward(loss)对应 engine.py 中的 backward()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_gradsFalse) else: loss.backward()关于梯度清零权重组完成一次 mini-batch 更新后DeepSpeed 会自动清空梯度无需像原生 PyTorch 那样显式调用optimizer.zero_grad()。此外DeepSpeed 在底层处理了分布式数据并行与 FP16这使得训练代码在多处得到简化(A) 梯度 allreduce 通信不再需要手动触发。DeepSpeed 会在梯度累积边界自动完成梯度平均因此 Megatron 原有的torch.distributed.all_reduce代码块要跳过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_afterFalse, fp32_allreduceargs.fp32_allreduce) timers(allreduce).stop()(B) master grads 更新与梯度裁剪一并交给 DeepSpeed。FP16 下的optimizer.update_master_grads()与clip_grad逻辑在 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()对应 engine.py 中的 step()会同时完成模型参数更新与学习率调度因此原生代码中的optimizer.step()lr_scheduler.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 1Loss Scaling 日志字段的替换GPT-2 训练脚本会在训练日志中打印 loss scaling 值。在 DeepSpeed 的 FP16 optimizer 内部该值存储在cur_scale字段中而不是 Megatron optimizer 的loss_scale。因此日志字符串需要做如下替换if args.fp16: log_string loss scale {:.1f} |.format( optimizer.cur_scale if args.deepspeed else optimizer.loss_scale)修改三Checkpoint 保存与加载对接DeepSpeed engine 提供了灵活的检查点 API能够同时管理「客户端模型自身的状态」与「DeepSpeed 内部状态」如优化器、lr scheduler、随机数状态等。对应 API 在 engine.py 中实现为def save_checkpoint(self, save_dir, tag, client_state{}) def load_checkpoint(self, load_dir, tag)要启用 DeepSpeed需要更新 Megatron 的utils.py中保存/加载检查点的逻辑。新增save_ds_checkpoint()并接入保存流程新建如下函数收集客户端模型状态iteration 与各类 RNG 状态再通过 engine 的save_checkpoint()一并写入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 原有的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: ......加载流程使用 DeepSpeed API在load_checkpoint()中使用 DeepSpeed 的加载接口并把客户端状态返回给上层。load_checkpoint()返回的元组中第一个元素是 checkpoint 名称若为None表示加载失败第二个元素即先前通过client_state存入的状态字典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: ......可选优化激活检查点Activation Checkpointing大模型训练中激活值activation的显存占用往往成为瓶颈。DeepSpeed 在模型并行训练中可以把激活检查点切分到各模型并行 GPU 上partition activation降低单卡激活显存或将激活检查点卸载到 CPU。这两类优化都是可选的仅当激活内存成为瓶颈时才需要启用。启用的方式是使用deepspeed.checkpointingAPI替换Megatron 的激活检查点与随机数状态追踪 API且替换必须发生在这些 API首次被调用之前。结合 checkpointing.py 中 configure() 的定义其支持partition_activations、contiguous_checkpointing、num_checkpoints、checkpoint_in_cpu、synchronize、profile等配置项。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_configargs.deepspeed_config, partition_activationTrue) 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_seedb) 在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对应的is_configured()见 checkpointing.py会返回真值各种激活检查点优化——激活切分partition_activations、连续内存检查点contiguous_checkpointing、CPU 检查点checkpoint_in_cpu——既可以在调用deepspeed.checkpointing.configure时以参数指定也可以在deepspeed_configJSON 配置文件中统一声明由args.deepspeed_config传入。启动 DeepSpeed 训练脚本假设webtext数据已按前述步骤准备完毕即可用如下脚本启动带 DeepSpeed 的 Megatron-LM GPT-2 训练脚本位于 Megatron 训练代码仓库内单 GPU 运行bash scripts/ds_pretrain_gpt.sh多 GPU / 多节点运行脚本名称即提示其使用 ZeRO-2 模型并行bash scripts/ds_zero2_pretrain_gpt_model_parallel.sh这两个脚本内部会以--deepspeed --deepspeed_config json的方式启动训练入口触发前文实现的 DeepSpeed 初始化与训练循环改造。结合 GPT-2 评估 ZeRO-1 / ZeRO-2 的规模化收益把 DeepSpeed 接入 Megatron 的意义最终要落到大规模训练的能力与效率上。该教程配套的官方评测描述如下ZeRO-12020 年 2 月发布执行 ZeRO 中的优化器状态切分ZeRO-22020 年 5 月扩展在 ZeRO-1 基础上增加梯度切分、激活切分与连续内存优化。ZeRO-2 大幅降低了训练超大模型的内存占用因此同样的模型可以用 i) 更低的模型并行度、ii) 更大的 batch size 来训练。更低的模型并行度会提升单个计算如矩阵乘法的粒度从而提高计算效率同时减少模型并行 GPU 之间的通信更大的 batch size 具有类似的提高计算粒度与减少通信的效果。这正是 ZeRO-2 与 Megatron 结合后「模型规模与速度双双提升」的根本原因。图ZeRO-2 可扩展到 1700 亿参数、吞吐最高提升约 10 倍、可获得超线性加速并在最高 130 亿参数规模内无需代码重构即可提升易用性图片与说明来自教程原文。该评测从四个维度总结了 DeepSpeed ZeRO-2 相对单独使用 Megatron 的优势模型规模当时主流大模型——OpenAI GPT-21.5B、NVIDIA Megatron-LM8.3B、Google T511B、Microsoft Turing-NLG17B——的参数量级为 1.5B 至 17BZeRO-2 则提供系统级支持可高效运行1700 亿参数的模型比上述最大模型高出一个数量级。训练速度ZeRO-2 可在 400 块 NVIDIA V100 GPU 集群上训练千亿参数模型单卡性能超过 38 TFlops、聚合性能超过 15 PFlops。同等规模模型下ZeRO-2 相对单独使用 Megatron 训练速度快约 10 倍相对 ZeRO-1 快约 5 倍。超线性扩展性观测到超线性加速——GPU 数量翻倍时性能提升超过一倍。这是因为数据并行度提高后 ZeRO-2 进一步摊薄了模型状态的内存占用使每卡可容纳更大 batch。降低大模型训练门槛ZeRO-2 让模型科学家在**不需要任何模型并行无需模型重构**的前提下高效训练高达 130 亿参数的模型——这一规模已超过当时多数最大模型如 110 亿参数的 T5。作为对比经典数据并行方案如 PyTorch DDP在 14 亿参数模型上即显存耗尽ZeRO-1 支持到约 60 亿参数。更进一步的评测还表明在无模型并行的场景下这些模型可以在低带宽集群上获得比使用模型并行明显更高的吞吐。例如在一个由 4 节点组成、通过 40 Gbps Infiniband 互联、每节点 4 块 NVIDIA 16GB V100PCI-E 连接的集群上GPT-2 使用 ZeRO 数据并行比使用模型并行快约 4 倍。换言之大模型训练不再被限制在拥有超高速互联的专用 GPU 集群中适度带宽的集群同样可行。需要说明的是以上数据是教程文档记录的 2020 年 ZeRO-2 发布评测对应 DeepSpeed 官方 ZeRO 技术博客与 ZeRO 论文结论涉及 NVIDIA V100 等当时主流硬件若需在具体软硬件组合上验证收益建议参照本仓库 docs/_tutorials/zero.md 与 ZeRO 配置说明 重新跑一轮基线对比。总结与延伸阅读把 DeepSpeed 接入 Megatron-LM GPT-2 的完整链路可概括为注册参数add_config_arguments→ 初始化引擎initialize传裸 Adam 与mpu→ 改造训练循环backward/step去掉手动 allreduce 与 master grads 管理→ 对接检查点save_checkpoint/load_checkpoint→ 可选激活切分deepspeed.checkpointing→ 通过ds_*脚本启动。所有关键 API 都可在本仓库核心代码中找到一一对应的实现参数注册与引擎初始化deepspeed/init.py 中的initialize()与 add_config_arguments()训练与检查点引擎deepspeed/runtime/engine.py 中的backwardL3288、stepL3500、save_checkpointL4844、load_checkpointL4366与is_gradient_accumulation_boundaryL3358激活检查点模块deepspeed/runtime/activation_checkpointing/checkpointing.py 中的configure()、is_configured()、checkpoint()等。想继续深入可以进一步阅读本仓库中的 Getting Started 教程、ZeRO 优化器教程、DeepSpeed 配置文件ds_config字段说明以及模型并行相关的 Megatron 框架原始论文与 ZeRO 论文从而把「能跑通」提升为「能调优、能规模化」。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 在 Megatron-LM GPT-2 上接入 DeepSpeed：模型并行训练的逐文件改造与 ZeRO-2 实践
在 Megatron-LM GPT-2 上接入 DeepSpeed模型并行训练的逐文件改造与 ZeRO-2 实践【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文是一份面向工程实战的技术指南完整讲解如何把 DeepSpeed 分布式训练优化库接入基于 Megatron-LM 的 GPT-2 大模型训练代码从数据与环境准备、原始模型跑通到逐文件改造参数解析、引擎初始化、前向/反向/更新循环与断点续训再到可选激活检查点优化与 ZeRO-2 规模化训练效果。读完本文你将能够复现官方教程中的完整接入流程并理解每一步背后 DeepSpeed 引擎Engine的源码级实现原理。本文内容以仓库内文档 docs/_tutorials/megatron.md 为骨架并结合本仓库中 DeepSpeed 核心源码展开印证。背景与适用范围什么样的模型适合走这条接入路径Megatron-LM 是一套支持**模型并行model parallelism与多节点multi-node**训练的 Transformer 训练框架其对应论文为Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism。GPT-2 是其中典型的超大规模 Transformer 模型。DeepSpeed 将 ZeRO 数据并行与 Megatron 的模型并行相结合后可以解决「模型大到单卡乃至单节点放不下」的训练难题。本教程面向的正是此类已具备模型并行能力、希望在不大规模改写逻辑的前提下叠加 ZeRO 优化的训练代码。需要先说明一个代码归属问题避免混淆本教程提到的megatron/arguments.py、pretrain_gpt.py、megatron/utils.py、mpu/transformer.py以及各类.sh启动脚本都位于Megatron-LM 训练代码仓库DeepSpeed 社区维护的 Megatron-DeepSpeed 分叉作为示例仓库的 submodule 引入而不是DeepSpeed 核心库所在目录中的文件。DeepSpeed 核心库即本仓库deepspeed/目录为接入过程提供的是deepspeed.initialize、deepspeed.add_config_arguments、Engine 的backward/step/save_checkpoint/load_checkpoint与激活检查点等 API。因此本仓库内并找不到megatron/目录需要先按下文拉取 Megatron 代码。如果你还没有阅读过入门指南建议先浏览 Getting Started 教程其中覆盖了 DeepSpeed 的安装、基础配置与最小训练示例。第一步用原始 Megatron-LM 跑通 GPT-2改造基线在接入 DeepSpeed 之前先在原始代码上跑通训练作为后续性能对比的基线。拉取 Megatron-LM 训练代码Megatron-LM 的原版模型代码已被复制到 DeepSpeed 的 Megatron 分叉中并以 submodule 形式提供。执行以下命令获取git submodule update --init --recursive训练数据准备下载 OpenAIwebtext数据集按 Megatron 官方 README 中Collecting GPT2 WebText Data一节提供的步骤下载数据并在DeepSpeedExamples/Megatron-LM/data目录下为数据放置一个符号链接使训练脚本可以按data路径找到语料。运行未修改的 Megatron-LM GPT-2单 GPU 运行修改examples/pretrain_gpt.sh把它的--train-data参数设为webtext然后执行bash examples/pretrain_gpt.sh多 GPU / 多节点运行修改examples/pretrain_gpt_distributed_with_mp.sh将--train-data参数设为webtextGPUS_PER_NODE表示测试中每个节点投入的 GPU 数量NNODES表示参与测试的节点数然后执行bash examples/pretrain_gpt_distributed_with_mp.sh启用 DeepSpeed 的总体改造清单使用 DeepSpeed 需要修改 Megatron-LM 训练代码中的三个文件分工如下文件职责megatron/arguments.py参数配置把 DeepSpeed 命令行参数注册进参数解析器pretrain_gpt.py训练主入口初始化 DeepSpeed 引擎并改造训练循环megatron/utils.py模型检查点的保存与加载逻辑对接 DeepSpeed修改一把 DeepSpeed 参数注入参数解析器DeepSpeed 提供了一个便捷函数deepspeed.add_config_arguments(parser)用于向已有的argparse.ArgumentParser注入一套 DeepSpeed 命令行参数。在megatron/arguments.py的get_args()末尾调用它def get_args(): Parse all the args. parser argparse.ArgumentParser(descriptionPyTorch 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)从源码看add_config_arguments()实际注入的核心参数包括--deepspeed布尔开关表示启用 DeepSpeed--deepspeed_config json 文件路径DeepSpeed 运行时 JSON 配置文件路径。正是因为这一处注册后续训练脚本里才可以通过--deepspeed与--deepspeed_config显式打开 DeepSpeed 并指定配置。修改二初始化 DeepSpeed 引擎并改造训练循环deepspeed.initialize是 DeepSpeed 引擎的入口它同时负责创建model engine、optimizer 与 LR scheduler。结合当前仓库 deepspeed/init.py 中的实际签名可见def initialize(args, model, optimizerNone, model_parametersNone, training_dataNone, lr_schedulerNone, mpuNone, dist_init_requiredTrue, collate_fnNone):model任何 wrapper 之前的原始nn.Moduleoptimizer用户自定义 optimizer若提供则覆盖 JSON 配置中的 optimizer 定义mpu提供get_{model,data}_parallel_{rank,group,world_size}()接口的模型并行单元对象用于把 Megatron 的模型并行信息告知 DeepSpeeddist_init_required是否由 DeepSpeed 自动初始化 torch 分布式环境。对于 Megatron-LM GPT-2在其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( modelmodel, optimizeroptimizer, argsargs, lr_schedulerlr_scheduler, mpumpu, dist_init_requiredFalse )注意这里dist_init_requiredFalseMegatron 训练代码本身已经完成了分布式环境的初始化含mpu建立的各种进程组因此需要明确告知 DeepSpeed 不要再重复初始化。一个重要细节FP16 下必须传「未包装」的 AdamMegatron-LM 在启用 FP16 时会给Adam优化器套一层自己的 FP16 wrapper负责维护 master weights / master grads 与动态 loss scaling。但DeepSpeed 拥有自己的 FP16 Optimizer因此当args.deepspeed为真时get_optimizer()必须返回未包装的裸Adam把 FP16 相关管理全部交给 DeepSpeeddef get_optimizer(model, args): Setup the optimizer. ...... # Use Adam. optimizer Adam(param_groups, lrargs.lr, weight_decayargs.weight_decay) if args.deepspeed: # fp16 wrapper is not required for DeepSpeed. return optimizer使用 Training API前向 / 反向 / 参数更新deepspeed.initialize返回的model是DeepSpeed Model EngineDeepSpeedEngine训练将围绕 engine 暴露的 forward、backward、step API 进行。前向传播前向传播 API 与 PyTorch 完全兼容不需要任何改动。反向传播反向传播直接调用 engine 的model.backward(loss)对应 engine.py 中的 backward()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_gradsFalse) else: loss.backward()关于梯度清零权重组完成一次 mini-batch 更新后DeepSpeed 会自动清空梯度无需像原生 PyTorch 那样显式调用optimizer.zero_grad()。此外DeepSpeed 在底层处理了分布式数据并行与 FP16这使得训练代码在多处得到简化(A) 梯度 allreduce 通信不再需要手动触发。DeepSpeed 会在梯度累积边界自动完成梯度平均因此 Megatron 原有的torch.distributed.all_reduce代码块要跳过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_afterFalse, fp32_allreduceargs.fp32_allreduce) timers(allreduce).stop()(B) master grads 更新与梯度裁剪一并交给 DeepSpeed。FP16 下的optimizer.update_master_grads()与clip_grad逻辑在 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()对应 engine.py 中的 step()会同时完成模型参数更新与学习率调度因此原生代码中的optimizer.step()lr_scheduler.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 1Loss Scaling 日志字段的替换GPT-2 训练脚本会在训练日志中打印 loss scaling 值。在 DeepSpeed 的 FP16 optimizer 内部该值存储在cur_scale字段中而不是 Megatron optimizer 的loss_scale。因此日志字符串需要做如下替换if args.fp16: log_string loss scale {:.1f} |.format( optimizer.cur_scale if args.deepspeed else optimizer.loss_scale)修改三Checkpoint 保存与加载对接DeepSpeed engine 提供了灵活的检查点 API能够同时管理「客户端模型自身的状态」与「DeepSpeed 内部状态」如优化器、lr scheduler、随机数状态等。对应 API 在 engine.py 中实现为def save_checkpoint(self, save_dir, tag, client_state{}) def load_checkpoint(self, load_dir, tag)要启用 DeepSpeed需要更新 Megatron 的utils.py中保存/加载检查点的逻辑。新增save_ds_checkpoint()并接入保存流程新建如下函数收集客户端模型状态iteration 与各类 RNG 状态再通过 engine 的save_checkpoint()一并写入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 原有的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: ......加载流程使用 DeepSpeed API在load_checkpoint()中使用 DeepSpeed 的加载接口并把客户端状态返回给上层。load_checkpoint()返回的元组中第一个元素是 checkpoint 名称若为None表示加载失败第二个元素即先前通过client_state存入的状态字典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: ......可选优化激活检查点Activation Checkpointing大模型训练中激活值activation的显存占用往往成为瓶颈。DeepSpeed 在模型并行训练中可以把激活检查点切分到各模型并行 GPU 上partition activation降低单卡激活显存或将激活检查点卸载到 CPU。这两类优化都是可选的仅当激活内存成为瓶颈时才需要启用。启用的方式是使用deepspeed.checkpointingAPI替换Megatron 的激活检查点与随机数状态追踪 API且替换必须发生在这些 API首次被调用之前。结合 checkpointing.py 中 configure() 的定义其支持partition_activations、contiguous_checkpointing、num_checkpoints、checkpoint_in_cpu、synchronize、profile等配置项。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_configargs.deepspeed_config, partition_activationTrue) 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_seedb) 在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对应的is_configured()见 checkpointing.py会返回真值各种激活检查点优化——激活切分partition_activations、连续内存检查点contiguous_checkpointing、CPU 检查点checkpoint_in_cpu——既可以在调用deepspeed.checkpointing.configure时以参数指定也可以在deepspeed_configJSON 配置文件中统一声明由args.deepspeed_config传入。启动 DeepSpeed 训练脚本假设webtext数据已按前述步骤准备完毕即可用如下脚本启动带 DeepSpeed 的 Megatron-LM GPT-2 训练脚本位于 Megatron 训练代码仓库内单 GPU 运行bash scripts/ds_pretrain_gpt.sh多 GPU / 多节点运行脚本名称即提示其使用 ZeRO-2 模型并行bash scripts/ds_zero2_pretrain_gpt_model_parallel.sh这两个脚本内部会以--deepspeed --deepspeed_config json的方式启动训练入口触发前文实现的 DeepSpeed 初始化与训练循环改造。结合 GPT-2 评估 ZeRO-1 / ZeRO-2 的规模化收益把 DeepSpeed 接入 Megatron 的意义最终要落到大规模训练的能力与效率上。该教程配套的官方评测描述如下ZeRO-12020 年 2 月发布执行 ZeRO 中的优化器状态切分ZeRO-22020 年 5 月扩展在 ZeRO-1 基础上增加梯度切分、激活切分与连续内存优化。ZeRO-2 大幅降低了训练超大模型的内存占用因此同样的模型可以用 i) 更低的模型并行度、ii) 更大的 batch size 来训练。更低的模型并行度会提升单个计算如矩阵乘法的粒度从而提高计算效率同时减少模型并行 GPU 之间的通信更大的 batch size 具有类似的提高计算粒度与减少通信的效果。这正是 ZeRO-2 与 Megatron 结合后「模型规模与速度双双提升」的根本原因。图ZeRO-2 可扩展到 1700 亿参数、吞吐最高提升约 10 倍、可获得超线性加速并在最高 130 亿参数规模内无需代码重构即可提升易用性图片与说明来自教程原文。该评测从四个维度总结了 DeepSpeed ZeRO-2 相对单独使用 Megatron 的优势模型规模当时主流大模型——OpenAI GPT-21.5B、NVIDIA Megatron-LM8.3B、Google T511B、Microsoft Turing-NLG17B——的参数量级为 1.5B 至 17BZeRO-2 则提供系统级支持可高效运行1700 亿参数的模型比上述最大模型高出一个数量级。训练速度ZeRO-2 可在 400 块 NVIDIA V100 GPU 集群上训练千亿参数模型单卡性能超过 38 TFlops、聚合性能超过 15 PFlops。同等规模模型下ZeRO-2 相对单独使用 Megatron 训练速度快约 10 倍相对 ZeRO-1 快约 5 倍。超线性扩展性观测到超线性加速——GPU 数量翻倍时性能提升超过一倍。这是因为数据并行度提高后 ZeRO-2 进一步摊薄了模型状态的内存占用使每卡可容纳更大 batch。降低大模型训练门槛ZeRO-2 让模型科学家在**不需要任何模型并行无需模型重构**的前提下高效训练高达 130 亿参数的模型——这一规模已超过当时多数最大模型如 110 亿参数的 T5。作为对比经典数据并行方案如 PyTorch DDP在 14 亿参数模型上即显存耗尽ZeRO-1 支持到约 60 亿参数。更进一步的评测还表明在无模型并行的场景下这些模型可以在低带宽集群上获得比使用模型并行明显更高的吞吐。例如在一个由 4 节点组成、通过 40 Gbps Infiniband 互联、每节点 4 块 NVIDIA 16GB V100PCI-E 连接的集群上GPT-2 使用 ZeRO 数据并行比使用模型并行快约 4 倍。换言之大模型训练不再被限制在拥有超高速互联的专用 GPU 集群中适度带宽的集群同样可行。需要说明的是以上数据是教程文档记录的 2020 年 ZeRO-2 发布评测对应 DeepSpeed 官方 ZeRO 技术博客与 ZeRO 论文结论涉及 NVIDIA V100 等当时主流硬件若需在具体软硬件组合上验证收益建议参照本仓库 docs/_tutorials/zero.md 与 ZeRO 配置说明 重新跑一轮基线对比。总结与延伸阅读把 DeepSpeed 接入 Megatron-LM GPT-2 的完整链路可概括为注册参数add_config_arguments→ 初始化引擎initialize传裸 Adam 与mpu→ 改造训练循环backward/step去掉手动 allreduce 与 master grads 管理→ 对接检查点save_checkpoint/load_checkpoint→ 可选激活切分deepspeed.checkpointing→ 通过ds_*脚本启动。所有关键 API 都可在本仓库核心代码中找到一一对应的实现参数注册与引擎初始化deepspeed/init.py 中的initialize()与 add_config_arguments()训练与检查点引擎deepspeed/runtime/engine.py 中的backwardL3288、stepL3500、save_checkpointL4844、load_checkpointL4366与is_gradient_accumulation_boundaryL3358激活检查点模块deepspeed/runtime/activation_checkpointing/checkpointing.py 中的configure()、is_configured()、checkpoint()等。想继续深入可以进一步阅读本仓库中的 Getting Started 教程、ZeRO 优化器教程、DeepSpeed 配置文件ds_config字段说明以及模型并行相关的 Megatron 框架原始论文与 ZeRO 论文从而把「能跑通」提升为「能调优、能规模化」。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考