尧图网站建设 
DeepSpeed BERT 预训练实战把 DeepSpeed 无缝接入现有 Transformer 预训练代码库【免费下载链接】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 官方的 BERT 预训练教程为主线完整演示如何把一个基于 HuggingFace Transformers 与 NVIDIA DeepLearningExamples 改造而来的原生 BERT 预训练代码Bing BERT 风格逐步接入 DeepSpeed从命令行参数注册、deepspeed.initialize()引擎初始化、backward()/step()训练循环改造到统一的检查点保存与恢复、JSON 配置编写再到可选的高性能 Transformer Kernel 与多机多卡启动方式。读完你不仅能照做把任何一个 BERT 类预训练工程跑在 DeepSpeed 上还能在当前仓库源码层面理解每个 API 背后的实现位置。教程背景两条 BERT 预训练路线与定位文档开头有一段需要认真对待的提示2022 年 8 月之后社区在配套的 Megatron-DeepSpeed 仓库中新增了一个基于 Pile 数据集的 BERT 预训练/微调示例位于其examples_deepspeed/bert_with_pile目录内含独立 README。相比本文下面讲解的经典示例新示例额外支持 ZeRO 与张量切片模型并行因此可支撑更大模型规模使用公开且更丰富的 Pile 数据集并参照相关论文对模型结构与超参数做了调整。官方建议如果目标是训练更大规模或更高质量的 BERT 风格模型优先参考 Megatron-DeepSpeed 中的新示例如果目标是严格复现原始 BERT 模型则跟随本教程基于 DeepSpeedExamples 仓库bing_bert目录的经典示例无论使用哪个示例本教程的核心价值在于讲清楚如何把 DeepSpeed 集成进一个预训练代码库——这一套接入方法论是通用的。本教程涉及的实际示例代码bing_bert目录及其中的模型配置文件、启动脚本位于配套的 DeepSpeedExamples 独立仓库不在当前 DeepSpeed 核心仓库内但教程中调用的每一个 DeepSpeed API都可以在当前核心仓库中找到对应实现这也是下文源码印证部分要做的事。不用 DeepSpeed 时先搭好数据管线并跑通原生 BERT教程给出的原始实现取材于 HuggingFace Transformers 与 NVIDIA DeepLearningExamples 两个上游项目在其基础上做了两处关键改造并托管于 DeepSpeedExamples 的bing_bert目录bing_bert/nvidia/采用 NVIDIA BERT 的建模modeling代码bing_bert/turing/扩展自微软 Project Turing 的数据管线代码。训练数据准备下载 Wikipedia 与 BookCorpus 数据集并把路径写入模型配置文件bing_bert/bert_large_adam_seq128.json{ datasets: { wiki_pretrain_dataset: /data/bert/bnorick_format/128/wiki_pretrain, bc_pretrain_dataset: /data/bert/bnorick_format/128/bookcorpus_pretrain } }两个数据集路径需要替换为真实绝对路径128表示按 128 token 序列长度切分好的预训练样本目录序列切分格式为 bnorickbing_bert 数据管线使用的二进制样本格式。文档亦注明下载与预处理的详细说明后续补充。原生启动命令在DeepSpeedExamples/bing_bert目录下先以未集成 DeepSpeed 的train.py跑通基线的数据加载与训练流程python train.py \ --cf bert_large_adam_seq128.json \ --train_batch_size 64 \ --max_seq_length 128 \ --gradient_accumulation_steps 1 \ --max_grad_norm 1.0 \ --fp16 \ --loss_scale 0 \ --delay_allreduce \ --max_steps 10 \ --output_dir path-to-model-output参数含义说明--cf指定模型配置文件BERT-Large、seq 128、Adam 优化器--train_batch_size训练 batch size--max_seq_length最大序列长度 128--gradient_accumulation_steps梯度累积步数基线为 1--max_grad_norm梯度裁剪阈值--fp16 --loss_scale 0启用 FP16 混合精度loss scale 为 0 表示使用动态 loss scaling--max_steps 10先跑 10 步做冒烟验证。启用 DeepSpeed 只需改动两个文件接入 DeepSpeed 时教程要求只编辑两个文件train.py训练主入口utils.py训练参数与检查点保存/加载的工具函数。加上一个新建的 DeepSpeed JSON 配置文件。下面分步展开每一步都标注当前核心仓库中的对应实现位置。第一步注册 DeepSpeed 命令行参数在train.py中调用deepspeed.add_config_arguments()把 DeepSpeed 专用参数挂到现有 argparse parser 上def get_arguments(): parser get_argument_parser() # Include DeepSpeed configuration arguments parser deepspeed.add_config_arguments(parser) args parser.parse_args() return args从源码看deepspeed/init.py 中的add_config_arguments()实际调用了内部助手_add_core_arguments()会向 parser 注入一个名为 DeepSpeed 的参数组核心新增参数包括--deepspeed布尔开关方便用户代码判断是否启用 DeepSpeed、--deepspeed_config jsonDeepSpeed JSON 配置文件路径以及一组用于向后兼容的旧名参数--deepscale/--deepscale_config源码会打印 deprecation 警告。也就是说加完这一行train.py就能识别 DeepSpeed 启动参数了。第二步用 deepspeed.initialize 构建模型引擎修改prepare_model_optimizer()把原始模型与优化器参数交给deepspeed.initialize()def prepare_model_optimizer(args): # Loading Model model BertMultiTask(args) # Optimizer parameters optimizer_parameters prepare_optimizer_parameters(args, model) model.network, optimizer, _, _ deepspeed.initialize(argsargs, modelmodel.network, model_parametersoptimizer_parameters, dist_init_requiredFalse) return model, optimizer需要注意Bing BERT 的原始模型被封装在model.network中因此这里传的是model.network而 DeepSpeed 返回的第一个值模型引擎被重新赋回model.network。之后model.network就由普通nn.Module变成了 DeepSpeed 引擎DeepSpeedEngine。对照源码deepspeed/init.py 的initialize()其完整签名依次接收args、model、optimizer、model_parameters、training_data、lr_scheduler、distributed_port、mpu、dist_init_required、collate_fn、config等参数并返回四元组(engine, optimizer, training_dataloader, lr_scheduler)若不传lr_scheduler与optimizer引擎会依据 JSON 配置自动构造若传了用户自建 optimizer/scheduler则引擎包装后返回dist_init_requiredFalse表示分布式初始化已由应用自行完成引擎不再重复初始化配置既可以通过args.deepspeed_config传入也可以通过initialize(config...)直接传字典或路径。initialize()内部会先初始化分布式通信随后读取配置并据此实例化DeepSpeedEngine对流水并行模型则实例化PipelineEngine启用 hybrid engine 配置时实例化DeepSpeedHybridEngine。第三步改造前向后的反向与优化器更新DeepSpeed 模型引擎暴露了与nn.Module相同的前向 API因此前向传播代码一行不用改只需替换反向传播与优化器更新。反向传播改为直接调用引擎的backward(loss)# Compute loss if args.deepspeed: model.network.backward(loss) else: if args.fp16: optimizer.backward(loss) else: loss.backward()参数更新改为调用引擎的step()梯度清零由 DeepSpeed 在每步权重更新后自动完成if args.deepspeed: model.network.step() else: optimizer.step() optimizer.zero_grad()从源码印证deepspeed/runtime/engine.pybackward(loss, retain_graphFalse, scale_wrt_gasTrue)在非托管梯度累积模式下backward()只做本地梯度累积accumulation 边界处的梯度规约与优化器更新统一在step()内触发同时会依据gradient_accumulation_steps()对 loss 进行缩放loss loss / gas并针对 ZeRO、AMP 等路径做 loss scaling 处理step(lr_kwargsNone)执行参数更新与学习率调整同时负责管理梯度清零与梯度累积边界逻辑引擎暴露的train_batch_size()、train_micro_batch_size_per_gpu()、gradient_accumulation_steps()属性同文件将自动把 JSON 里的批大小配置解析为训练循环可用的语义开发者无需再手工维护这些数值。这也解释了为什么改了这两处就能跑DeepSpeed 把梯度累积、loss 缩放、梯度裁剪、梯度清零、数据并行通信等分布式训练样板逻辑全部收敛进了引擎内部。第四步用统一检查点 API 保存/恢复训练状态DeepSpeed 模型引擎的检查点 API 同时管理客户端模型状态与DeepSpeed 自身内部状态如优化器状态、梯度累积步数、ZeRO 分区后的状态等签名如下def save_checkpoint(self, save_dir, tag, client_state{}) def load_checkpoint(self, load_dir, tag)在train.py的checkpoint_model()中收集客户端状态并交给引擎保存def checkpoint_model(PATH, ckpt_id, model, epoch, last_global_step, last_global_data_samples, **kwargs): Utility function for checkpointing model optimizer dictionaries The main purpose for this is to be able to resume training from that instant again checkpoint_state_dict {epoch: epoch, last_global_step: last_global_step, last_global_data_samples: last_global_data_samples} # Add extra kwargs too checkpoint_state_dict.update(kwargs) success model.network.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict) return在load_training_checkpoint()中使用加载 API并把客户端状态取回def load_training_checkpoint(args, model, PATH, ckpt_id): Utility function for checkpointing model optimizer dictionaries The main purpose for this is to be able to resume training from that instant again _, checkpoint_state_dict model.network.load_checkpoint(PATH, ckpt_id) epoch checkpoint_state_dict[epoch] last_global_step checkpoint_state_dict[last_global_step] last_global_data_samples checkpoint_state_dict[last_global_data_samples] del checkpoint_state_dict return (epoch, last_global_step, last_global_data_samples)实现位置同样在 deepspeed/runtime/engine.pyload_checkpoint()位于第 4366 行附近、save_checkpoint()位于第 4844 行附近。保存时以(save_dir, tag)组织检查点目录与标签client_state以 dict 形式与模型/优化器状态一同落盘加载时返回(加载到的检查点路径, client_state)调用方再从返回的 dict 中取出自己写入的 epoch、全局步数、全局样本数等信息用于恢复训练进度。这套 API 也是后续教程如 ZeRO、通用检查点在更大模型上断点续训的基础。DeepSpeed JSON 配置文件批大小、优化器与 FP16 的声明式描述接入的最后一步是新建一个 DeepSpeed 配置 JSON例如deepspeed_bsz4096_adam_config.json。该文件集中声明批大小、优化器与参数、是否启用 FP16 等 DeepSpeed 专属配置{ train_batch_size: 4096, train_micro_batch_size_per_gpu: 64, steps_per_print: 1000, optimizer: { type: Adam, params: { lr: 2e-4, max_grad_norm: 1.0, weight_decay: 0.01, bias_correction: false } }, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 } }配置要点逐条说明train_batch_size: 4096期望的有效全局批大小等效 batch size为 4096train_micro_batch_size_per_gpu: 64单卡显存可即时容纳的 micro batch size 为 64DeepSpeed 会在引擎内部据此自动推导梯度累积步数gradient_accumulation_steps train_batch_size / (micro_batch_size_per_gpu × data_parallel_world_size)因此训练脚本里无需手工计算——这也正是引擎里gradient_accumulation_steps()属性存在的原因optimizer使用 Adam 优化器并给出学习率lr2e-4、梯度裁剪阈值max_grad_norm1.0、权重衰减weight_decay0.01bias_correction: false表示关闭 Adam 的偏置修正典型用于 FP16 大 batch 训练fp16启用 FP16 混合精度训练。loss_scale: 0表示使用动态 loss scalinginitial_scale_power: 16表示动态 loss scaling 的初始缩放因子为2^16。源码印证配置中的这些键会被DeepSpeedConfigdeepspeed/runtime/config.py解析FP16 配置最终落到引擎的dynamic_loss_scale()/dynamic_loss_scale_args()相关逻辑deepspeed/runtime/engine.pysteps_per_print: 1000则控制日志打印频率。值得注意的是配置里的 Adam 类型名做了大小写归一化源码中优化器名如ADAM_OPTIMIZER adam、ADAMW_OPTIMIZER、LAMB_OPTIMIZER等定义于 deepspeed/runtime/config.py在 JSON 中写Adam、adam均可被识别。到这里代码侧的改造就已经全部完成了。文档特别说明DeepSpeedExamples 的bing_bert目录内提供了一个已应用上述全部改动的deepspeed_train.py可直接对照。可选开启 DeepSpeed Transformer Kernel 加速为获得更高性能可以启用 DeepSpeed 的 Transformer Kernel融合 Transformer 层的专用算子。步骤分为两处在 utils.py 增加开关参数parser.add_argument(--deepspeed_transformer_kernel, defaultFalse, actionstore_true, helpUse DeepSpeed transformer kernel to accelerate.)默认False便于随时开关。在 BertEncoder 中替换 Transformer 层在建模源码的BertEncoder类中根据开关用 DeepSpeed transformer kernel 实例化层if args.deepspeed_transformer_kernel: from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig, DeepSpeedConfig if hasattr(args, deepspeed_config) and args.deepspeed_config: ds_config DeepSpeedConfig(args.deepspeed_config) else: raise RuntimeError(deepspeed_config is not found in args.) cuda_config DeepSpeedTransformerConfig( batch_size ds_config.train_micro_batch_size_per_gpu, max_seq_length args.max_seq_length, hidden_size config.hidden_size, heads config.num_attention_heads, attn_dropout_ratio config.attention_probs_dropout_prob, hidden_dropout_ratio config.hidden_dropout_prob, num_hidden_layers config.num_hidden_layers, initializer_range config.initializer_range, local_rank args.local_rank if hasattr(args, local_rank) else -1, seed args.seed, fp16 ds_config.fp16_enabled, pre_layer_normTrue, attn_dropout_checkpointargs.attention_dropout_checkpoint, normalize_invertibleargs.normalize_invertible, gelu_checkpointargs.gelu_checkpoint, stochastic_modeTrue) layer DeepSpeedTransformerLayer(cuda_config) else: layer BertLayer(config) self.layer nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])注意此时模型需要能访问到args因为 kernel 的绝大多数配置都来自 DeepSpeed 配置文件与命令行参数例如从ds_config读取 micro batch size 与是否启用 FP16。源码印证DeepSpeedTransformerLayer、DeepSpeedTransformerConfig与DeepSpeedConfig至今仍从deepspeed顶层导出见 deepspeed/init.py 中对deepspeed.ops.transformer与deepspeed.runtime.config的导入DeepSpeedTransformerConfig的完整字段定义在 deepspeed/ops/transformer/transformer.py其 Docstring 对stochastic_mode、normalize_invertible、gelu_checkpoint、pre_layer_norm等开关的作用与默认值均有明确说明可与教程代码互相印证。更多原理性内容可阅读同仓库教程 transformer_kernel.md。四条重要使用注意事项原样继承自教程batch_size是输入数据的最大 batch size所有微调/预测数据的 batch 都不能超过该阈值否则会抛异常。在 DeepSpeed 配置中 micro batch size 对应train_micro_batch_size_per_gpu。例如配置为 8、而预测要用 batch 12可以把 kernel 的 batch size 设为 12或用--predict_batch_size把预测 batch 降到 8 或更小。local_rank用于把 kernel 分配到正确设备。由于模型在此前已经执行过set_device()因此这里不设置通常也没有问题。stochastic_mode开启后性能更高但带有一定非确定性。预训练开启、微调关闭微调需要可复现性此点与 deepspeed/ops/transformer/transformer.py 中对stochastic_mode的建议一致。Transformer kernel 拥有自己专属的参数布局kernel 模式产出的检查点必须由同样开启了 kernel 的模型加载例如用于微调时。多节点启动4 节点 × 4 卡训练示例用deepspeed启动器运行deepspeed_train.py示例为四个节点、每节点四卡deepspeed --num_nodes 4 \ deepspeed_train.py \ --deepspeed \ --deepspeed_config deepspeed_bsz4096_adam_config.json \ --cf /path-to-deepspeed/examples/tests/bing_bert/bert_large_adam_seq128.json \ --train_batch_size 4096 \ --max_seq_length 128 \ --gradient_accumulation_steps 4 \ --max_grad_norm 1.0 \ --fp16 \ --loss_scale 0 \ --delay_allreduce \ --max_steps 32 \ --print_steps 1 \ --deepspeed_transformer_kernel \ --output_dir output_directory参数协同关系解读--deepspeed与--deepspeed_config启用 DeepSpeed 并指向上文创建的 JSON 配置--cf模型配置文件BERT-Large/seq128/Adam命令行中的--train_batch_size 4096与配置文件中保持一致配合单卡 micro batch 64、16 张卡与--gradient_accumulation_steps 4即可凑出 4096 的有效全局批大小64×16×4 4096--deepspeed_transformer_kernel打开 Transformer Kernel 加速每个节点内部的多卡数如 4通过启动器自动发现更通用的启动参数与注意事项参见仓库教程 getting-started.md。复现教程记载的最快 BERT 训练规模的配置教程记录于其发布时代宣称在保持 SQuAD 1.1 dev 集 F1 ≥ 90.5 竞争力的前提下取得了当时最快的 BERT 训练时间。文中给出的对比结果如下以下数字均为该教程文档当时记录的公开结果复现所需脚本与 JSON 在配套 DeepSpeedExamples 仓库的bing_bert目录下可参考其中的ds_train_bert_bsz64k_seq128.sh与ds_train_bert_bsz32k_seq512.sh使用 1024 张 V10064 个 NVIDIA DGX-2 节点在44 分钟内完成 BERT 预训练文档记载彼时对比的 NVIDIA 结果需 1472 张 V100、耗时 47 分钟即 DeepSpeed 不仅更快且少用约 30% 资源相比 Google 原始 BERT 在 64 颗 TPU2 上约 96 小时达到同等水平教程记载其 4 个 DGX-2 节点64 张 V100上训练时间少于 9 小时256 张 GPU 上耗时 2.4 小时快于文档记载的 NVIDIA 同卡数参考值 3.9 小时。不同规模下的训练耗时教程记载微调验证流程见 bert-finetuning.md节点数V100 GPU 数训练耗时1 DGX-21633 hr 13 min4 DGX-2648 hr 41 min16 DGX-2256144 min64 DGX-2102444 min用于复现上述结果的训练配置摘要如下详细脚本与配置位于配套 DeepSpeedExamples 仓库bing_bert目录参数128 序列长度512 序列长度总 batch size64K32K单卡 micro batch size648优化器LambLamb学习率11e-32e-3初始学习率lr_offset10e-40.0Min Lamb 系数0.010.01Max Lamb 系数0.30.3LR 调度器warmup_exp_decay_expwarmup_exp_decay_expWarmup 比例0.020.02Decay 率0.900.90Decay 步数250150最大训练步数75007500Rewarm 学习率N/ATrue输出检查点数150160-162样本数403M18-22MEpoch 数150160-162可复现性的两个关键点seq 128 使用 64K 的总 batch 配合 Lamb 优化器与 11e-3 高学习率seq 512 场景因上下文更长而把 batch 降为 32K、学习率降为 2e-3并启用了 rewarm。若要与预训练产出模型衔接下游任务验证请按 bert-finetuning.md 教程对 kernel 预训练模型做微调复现 SQuAD F1 指标。单卡吞吐结果Transformer Kernel 的收益可视化与上述大规模数据并行结果相辅相成的是单卡性能。下图文档原图存放于 docs/assets/images 目录展示了在 seq 128 与 seq 512 两种序列长度下经 DeepSpeed 优化的 BERT-Large 训练单卡吞吐对比教程记载相比 NVIDIA BERT 与 HuggingFace BERT 两个当时主流的 PyTorch 实现DeepSpeed 在 seq 128 与 seq 512 下分别达到约 64 与 53 teraflops 的单卡吞吐对应约 272 与 52 samples/second相对 NVIDIA BERT 提升最高约 28%、相对 HuggingFace BERT 提升最高约 62%同时支持在显存不溢出前提下把单卡 batch size 放大至约 1.8 倍。这些单卡能力的来源正是 Transformer Kernel 对多头注意力、LayerNorm、GELU 等算子的融合与算子级内存优化——相关配置字段与开关如pre_layer_norm、normalize_invertible、gelu_checkpoint、stochastic_mode的具体语义可在上一节及 transformer_kernel.md、transformer.py 中进一步研读。小结一套可以复用到任意 BERT 类项目的接入路径回看整个教程把 DeepSpeed 集成进 BERT 预训练代码库的完整动作可以浓缩为五步(1)用add_config_arguments()让train.py认识 DeepSpeed 参数(2)用deepspeed.initialize()把模型/优化器包成 DeepSpeed 引擎(3)把训练循环中的反向与参数更新换成引擎的backward()与step()(4)用save_checkpoint()/load_checkpoint()统一保存与恢复客户端状态和引擎内部状态(5)写一个声明批大小、优化器与 FP16 策略的 JSON 配置文件。再加上可选的 Transformer Kernel 融合层即可在同一份代码上获得从单卡吞吐优化到千卡级扩展的完整能力。这套方法与具体 BERT 实现NVIDIA/HuggingFace/Megatron-DeepSpeed解耦可平滑迁移到其他 Transformer 预训练工程。【免费下载链接】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），仅供参考 DeepSpeed BERT 预训练实战：把 DeepSpeed 无缝接入现有 Transformer 预训练代码库
DeepSpeed BERT 预训练实战把 DeepSpeed 无缝接入现有 Transformer 预训练代码库【免费下载链接】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 官方的 BERT 预训练教程为主线完整演示如何把一个基于 HuggingFace Transformers 与 NVIDIA DeepLearningExamples 改造而来的原生 BERT 预训练代码Bing BERT 风格逐步接入 DeepSpeed从命令行参数注册、deepspeed.initialize()引擎初始化、backward()/step()训练循环改造到统一的检查点保存与恢复、JSON 配置编写再到可选的高性能 Transformer Kernel 与多机多卡启动方式。读完你不仅能照做把任何一个 BERT 类预训练工程跑在 DeepSpeed 上还能在当前仓库源码层面理解每个 API 背后的实现位置。教程背景两条 BERT 预训练路线与定位文档开头有一段需要认真对待的提示2022 年 8 月之后社区在配套的 Megatron-DeepSpeed 仓库中新增了一个基于 Pile 数据集的 BERT 预训练/微调示例位于其examples_deepspeed/bert_with_pile目录内含独立 README。相比本文下面讲解的经典示例新示例额外支持 ZeRO 与张量切片模型并行因此可支撑更大模型规模使用公开且更丰富的 Pile 数据集并参照相关论文对模型结构与超参数做了调整。官方建议如果目标是训练更大规模或更高质量的 BERT 风格模型优先参考 Megatron-DeepSpeed 中的新示例如果目标是严格复现原始 BERT 模型则跟随本教程基于 DeepSpeedExamples 仓库bing_bert目录的经典示例无论使用哪个示例本教程的核心价值在于讲清楚如何把 DeepSpeed 集成进一个预训练代码库——这一套接入方法论是通用的。本教程涉及的实际示例代码bing_bert目录及其中的模型配置文件、启动脚本位于配套的 DeepSpeedExamples 独立仓库不在当前 DeepSpeed 核心仓库内但教程中调用的每一个 DeepSpeed API都可以在当前核心仓库中找到对应实现这也是下文源码印证部分要做的事。不用 DeepSpeed 时先搭好数据管线并跑通原生 BERT教程给出的原始实现取材于 HuggingFace Transformers 与 NVIDIA DeepLearningExamples 两个上游项目在其基础上做了两处关键改造并托管于 DeepSpeedExamples 的bing_bert目录bing_bert/nvidia/采用 NVIDIA BERT 的建模modeling代码bing_bert/turing/扩展自微软 Project Turing 的数据管线代码。训练数据准备下载 Wikipedia 与 BookCorpus 数据集并把路径写入模型配置文件bing_bert/bert_large_adam_seq128.json{ datasets: { wiki_pretrain_dataset: /data/bert/bnorick_format/128/wiki_pretrain, bc_pretrain_dataset: /data/bert/bnorick_format/128/bookcorpus_pretrain } }两个数据集路径需要替换为真实绝对路径128表示按 128 token 序列长度切分好的预训练样本目录序列切分格式为 bnorickbing_bert 数据管线使用的二进制样本格式。文档亦注明下载与预处理的详细说明后续补充。原生启动命令在DeepSpeedExamples/bing_bert目录下先以未集成 DeepSpeed 的train.py跑通基线的数据加载与训练流程python train.py \ --cf bert_large_adam_seq128.json \ --train_batch_size 64 \ --max_seq_length 128 \ --gradient_accumulation_steps 1 \ --max_grad_norm 1.0 \ --fp16 \ --loss_scale 0 \ --delay_allreduce \ --max_steps 10 \ --output_dir path-to-model-output参数含义说明--cf指定模型配置文件BERT-Large、seq 128、Adam 优化器--train_batch_size训练 batch size--max_seq_length最大序列长度 128--gradient_accumulation_steps梯度累积步数基线为 1--max_grad_norm梯度裁剪阈值--fp16 --loss_scale 0启用 FP16 混合精度loss scale 为 0 表示使用动态 loss scaling--max_steps 10先跑 10 步做冒烟验证。启用 DeepSpeed 只需改动两个文件接入 DeepSpeed 时教程要求只编辑两个文件train.py训练主入口utils.py训练参数与检查点保存/加载的工具函数。加上一个新建的 DeepSpeed JSON 配置文件。下面分步展开每一步都标注当前核心仓库中的对应实现位置。第一步注册 DeepSpeed 命令行参数在train.py中调用deepspeed.add_config_arguments()把 DeepSpeed 专用参数挂到现有 argparse parser 上def get_arguments(): parser get_argument_parser() # Include DeepSpeed configuration arguments parser deepspeed.add_config_arguments(parser) args parser.parse_args() return args从源码看deepspeed/init.py 中的add_config_arguments()实际调用了内部助手_add_core_arguments()会向 parser 注入一个名为 DeepSpeed 的参数组核心新增参数包括--deepspeed布尔开关方便用户代码判断是否启用 DeepSpeed、--deepspeed_config jsonDeepSpeed JSON 配置文件路径以及一组用于向后兼容的旧名参数--deepscale/--deepscale_config源码会打印 deprecation 警告。也就是说加完这一行train.py就能识别 DeepSpeed 启动参数了。第二步用 deepspeed.initialize 构建模型引擎修改prepare_model_optimizer()把原始模型与优化器参数交给deepspeed.initialize()def prepare_model_optimizer(args): # Loading Model model BertMultiTask(args) # Optimizer parameters optimizer_parameters prepare_optimizer_parameters(args, model) model.network, optimizer, _, _ deepspeed.initialize(argsargs, modelmodel.network, model_parametersoptimizer_parameters, dist_init_requiredFalse) return model, optimizer需要注意Bing BERT 的原始模型被封装在model.network中因此这里传的是model.network而 DeepSpeed 返回的第一个值模型引擎被重新赋回model.network。之后model.network就由普通nn.Module变成了 DeepSpeed 引擎DeepSpeedEngine。对照源码deepspeed/init.py 的initialize()其完整签名依次接收args、model、optimizer、model_parameters、training_data、lr_scheduler、distributed_port、mpu、dist_init_required、collate_fn、config等参数并返回四元组(engine, optimizer, training_dataloader, lr_scheduler)若不传lr_scheduler与optimizer引擎会依据 JSON 配置自动构造若传了用户自建 optimizer/scheduler则引擎包装后返回dist_init_requiredFalse表示分布式初始化已由应用自行完成引擎不再重复初始化配置既可以通过args.deepspeed_config传入也可以通过initialize(config...)直接传字典或路径。initialize()内部会先初始化分布式通信随后读取配置并据此实例化DeepSpeedEngine对流水并行模型则实例化PipelineEngine启用 hybrid engine 配置时实例化DeepSpeedHybridEngine。第三步改造前向后的反向与优化器更新DeepSpeed 模型引擎暴露了与nn.Module相同的前向 API因此前向传播代码一行不用改只需替换反向传播与优化器更新。反向传播改为直接调用引擎的backward(loss)# Compute loss if args.deepspeed: model.network.backward(loss) else: if args.fp16: optimizer.backward(loss) else: loss.backward()参数更新改为调用引擎的step()梯度清零由 DeepSpeed 在每步权重更新后自动完成if args.deepspeed: model.network.step() else: optimizer.step() optimizer.zero_grad()从源码印证deepspeed/runtime/engine.pybackward(loss, retain_graphFalse, scale_wrt_gasTrue)在非托管梯度累积模式下backward()只做本地梯度累积accumulation 边界处的梯度规约与优化器更新统一在step()内触发同时会依据gradient_accumulation_steps()对 loss 进行缩放loss loss / gas并针对 ZeRO、AMP 等路径做 loss scaling 处理step(lr_kwargsNone)执行参数更新与学习率调整同时负责管理梯度清零与梯度累积边界逻辑引擎暴露的train_batch_size()、train_micro_batch_size_per_gpu()、gradient_accumulation_steps()属性同文件将自动把 JSON 里的批大小配置解析为训练循环可用的语义开发者无需再手工维护这些数值。这也解释了为什么改了这两处就能跑DeepSpeed 把梯度累积、loss 缩放、梯度裁剪、梯度清零、数据并行通信等分布式训练样板逻辑全部收敛进了引擎内部。第四步用统一检查点 API 保存/恢复训练状态DeepSpeed 模型引擎的检查点 API 同时管理客户端模型状态与DeepSpeed 自身内部状态如优化器状态、梯度累积步数、ZeRO 分区后的状态等签名如下def save_checkpoint(self, save_dir, tag, client_state{}) def load_checkpoint(self, load_dir, tag)在train.py的checkpoint_model()中收集客户端状态并交给引擎保存def checkpoint_model(PATH, ckpt_id, model, epoch, last_global_step, last_global_data_samples, **kwargs): Utility function for checkpointing model optimizer dictionaries The main purpose for this is to be able to resume training from that instant again checkpoint_state_dict {epoch: epoch, last_global_step: last_global_step, last_global_data_samples: last_global_data_samples} # Add extra kwargs too checkpoint_state_dict.update(kwargs) success model.network.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict) return在load_training_checkpoint()中使用加载 API并把客户端状态取回def load_training_checkpoint(args, model, PATH, ckpt_id): Utility function for checkpointing model optimizer dictionaries The main purpose for this is to be able to resume training from that instant again _, checkpoint_state_dict model.network.load_checkpoint(PATH, ckpt_id) epoch checkpoint_state_dict[epoch] last_global_step checkpoint_state_dict[last_global_step] last_global_data_samples checkpoint_state_dict[last_global_data_samples] del checkpoint_state_dict return (epoch, last_global_step, last_global_data_samples)实现位置同样在 deepspeed/runtime/engine.pyload_checkpoint()位于第 4366 行附近、save_checkpoint()位于第 4844 行附近。保存时以(save_dir, tag)组织检查点目录与标签client_state以 dict 形式与模型/优化器状态一同落盘加载时返回(加载到的检查点路径, client_state)调用方再从返回的 dict 中取出自己写入的 epoch、全局步数、全局样本数等信息用于恢复训练进度。这套 API 也是后续教程如 ZeRO、通用检查点在更大模型上断点续训的基础。DeepSpeed JSON 配置文件批大小、优化器与 FP16 的声明式描述接入的最后一步是新建一个 DeepSpeed 配置 JSON例如deepspeed_bsz4096_adam_config.json。该文件集中声明批大小、优化器与参数、是否启用 FP16 等 DeepSpeed 专属配置{ train_batch_size: 4096, train_micro_batch_size_per_gpu: 64, steps_per_print: 1000, optimizer: { type: Adam, params: { lr: 2e-4, max_grad_norm: 1.0, weight_decay: 0.01, bias_correction: false } }, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 } }配置要点逐条说明train_batch_size: 4096期望的有效全局批大小等效 batch size为 4096train_micro_batch_size_per_gpu: 64单卡显存可即时容纳的 micro batch size 为 64DeepSpeed 会在引擎内部据此自动推导梯度累积步数gradient_accumulation_steps train_batch_size / (micro_batch_size_per_gpu × data_parallel_world_size)因此训练脚本里无需手工计算——这也正是引擎里gradient_accumulation_steps()属性存在的原因optimizer使用 Adam 优化器并给出学习率lr2e-4、梯度裁剪阈值max_grad_norm1.0、权重衰减weight_decay0.01bias_correction: false表示关闭 Adam 的偏置修正典型用于 FP16 大 batch 训练fp16启用 FP16 混合精度训练。loss_scale: 0表示使用动态 loss scalinginitial_scale_power: 16表示动态 loss scaling 的初始缩放因子为2^16。源码印证配置中的这些键会被DeepSpeedConfigdeepspeed/runtime/config.py解析FP16 配置最终落到引擎的dynamic_loss_scale()/dynamic_loss_scale_args()相关逻辑deepspeed/runtime/engine.pysteps_per_print: 1000则控制日志打印频率。值得注意的是配置里的 Adam 类型名做了大小写归一化源码中优化器名如ADAM_OPTIMIZER adam、ADAMW_OPTIMIZER、LAMB_OPTIMIZER等定义于 deepspeed/runtime/config.py在 JSON 中写Adam、adam均可被识别。到这里代码侧的改造就已经全部完成了。文档特别说明DeepSpeedExamples 的bing_bert目录内提供了一个已应用上述全部改动的deepspeed_train.py可直接对照。可选开启 DeepSpeed Transformer Kernel 加速为获得更高性能可以启用 DeepSpeed 的 Transformer Kernel融合 Transformer 层的专用算子。步骤分为两处在 utils.py 增加开关参数parser.add_argument(--deepspeed_transformer_kernel, defaultFalse, actionstore_true, helpUse DeepSpeed transformer kernel to accelerate.)默认False便于随时开关。在 BertEncoder 中替换 Transformer 层在建模源码的BertEncoder类中根据开关用 DeepSpeed transformer kernel 实例化层if args.deepspeed_transformer_kernel: from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig, DeepSpeedConfig if hasattr(args, deepspeed_config) and args.deepspeed_config: ds_config DeepSpeedConfig(args.deepspeed_config) else: raise RuntimeError(deepspeed_config is not found in args.) cuda_config DeepSpeedTransformerConfig( batch_size ds_config.train_micro_batch_size_per_gpu, max_seq_length args.max_seq_length, hidden_size config.hidden_size, heads config.num_attention_heads, attn_dropout_ratio config.attention_probs_dropout_prob, hidden_dropout_ratio config.hidden_dropout_prob, num_hidden_layers config.num_hidden_layers, initializer_range config.initializer_range, local_rank args.local_rank if hasattr(args, local_rank) else -1, seed args.seed, fp16 ds_config.fp16_enabled, pre_layer_normTrue, attn_dropout_checkpointargs.attention_dropout_checkpoint, normalize_invertibleargs.normalize_invertible, gelu_checkpointargs.gelu_checkpoint, stochastic_modeTrue) layer DeepSpeedTransformerLayer(cuda_config) else: layer BertLayer(config) self.layer nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])注意此时模型需要能访问到args因为 kernel 的绝大多数配置都来自 DeepSpeed 配置文件与命令行参数例如从ds_config读取 micro batch size 与是否启用 FP16。源码印证DeepSpeedTransformerLayer、DeepSpeedTransformerConfig与DeepSpeedConfig至今仍从deepspeed顶层导出见 deepspeed/init.py 中对deepspeed.ops.transformer与deepspeed.runtime.config的导入DeepSpeedTransformerConfig的完整字段定义在 deepspeed/ops/transformer/transformer.py其 Docstring 对stochastic_mode、normalize_invertible、gelu_checkpoint、pre_layer_norm等开关的作用与默认值均有明确说明可与教程代码互相印证。更多原理性内容可阅读同仓库教程 transformer_kernel.md。四条重要使用注意事项原样继承自教程batch_size是输入数据的最大 batch size所有微调/预测数据的 batch 都不能超过该阈值否则会抛异常。在 DeepSpeed 配置中 micro batch size 对应train_micro_batch_size_per_gpu。例如配置为 8、而预测要用 batch 12可以把 kernel 的 batch size 设为 12或用--predict_batch_size把预测 batch 降到 8 或更小。local_rank用于把 kernel 分配到正确设备。由于模型在此前已经执行过set_device()因此这里不设置通常也没有问题。stochastic_mode开启后性能更高但带有一定非确定性。预训练开启、微调关闭微调需要可复现性此点与 deepspeed/ops/transformer/transformer.py 中对stochastic_mode的建议一致。Transformer kernel 拥有自己专属的参数布局kernel 模式产出的检查点必须由同样开启了 kernel 的模型加载例如用于微调时。多节点启动4 节点 × 4 卡训练示例用deepspeed启动器运行deepspeed_train.py示例为四个节点、每节点四卡deepspeed --num_nodes 4 \ deepspeed_train.py \ --deepspeed \ --deepspeed_config deepspeed_bsz4096_adam_config.json \ --cf /path-to-deepspeed/examples/tests/bing_bert/bert_large_adam_seq128.json \ --train_batch_size 4096 \ --max_seq_length 128 \ --gradient_accumulation_steps 4 \ --max_grad_norm 1.0 \ --fp16 \ --loss_scale 0 \ --delay_allreduce \ --max_steps 32 \ --print_steps 1 \ --deepspeed_transformer_kernel \ --output_dir output_directory参数协同关系解读--deepspeed与--deepspeed_config启用 DeepSpeed 并指向上文创建的 JSON 配置--cf模型配置文件BERT-Large/seq128/Adam命令行中的--train_batch_size 4096与配置文件中保持一致配合单卡 micro batch 64、16 张卡与--gradient_accumulation_steps 4即可凑出 4096 的有效全局批大小64×16×4 4096--deepspeed_transformer_kernel打开 Transformer Kernel 加速每个节点内部的多卡数如 4通过启动器自动发现更通用的启动参数与注意事项参见仓库教程 getting-started.md。复现教程记载的最快 BERT 训练规模的配置教程记录于其发布时代宣称在保持 SQuAD 1.1 dev 集 F1 ≥ 90.5 竞争力的前提下取得了当时最快的 BERT 训练时间。文中给出的对比结果如下以下数字均为该教程文档当时记录的公开结果复现所需脚本与 JSON 在配套 DeepSpeedExamples 仓库的bing_bert目录下可参考其中的ds_train_bert_bsz64k_seq128.sh与ds_train_bert_bsz32k_seq512.sh使用 1024 张 V10064 个 NVIDIA DGX-2 节点在44 分钟内完成 BERT 预训练文档记载彼时对比的 NVIDIA 结果需 1472 张 V100、耗时 47 分钟即 DeepSpeed 不仅更快且少用约 30% 资源相比 Google 原始 BERT 在 64 颗 TPU2 上约 96 小时达到同等水平教程记载其 4 个 DGX-2 节点64 张 V100上训练时间少于 9 小时256 张 GPU 上耗时 2.4 小时快于文档记载的 NVIDIA 同卡数参考值 3.9 小时。不同规模下的训练耗时教程记载微调验证流程见 bert-finetuning.md节点数V100 GPU 数训练耗时1 DGX-21633 hr 13 min4 DGX-2648 hr 41 min16 DGX-2256144 min64 DGX-2102444 min用于复现上述结果的训练配置摘要如下详细脚本与配置位于配套 DeepSpeedExamples 仓库bing_bert目录参数128 序列长度512 序列长度总 batch size64K32K单卡 micro batch size648优化器LambLamb学习率11e-32e-3初始学习率lr_offset10e-40.0Min Lamb 系数0.010.01Max Lamb 系数0.30.3LR 调度器warmup_exp_decay_expwarmup_exp_decay_expWarmup 比例0.020.02Decay 率0.900.90Decay 步数250150最大训练步数75007500Rewarm 学习率N/ATrue输出检查点数150160-162样本数403M18-22MEpoch 数150160-162可复现性的两个关键点seq 128 使用 64K 的总 batch 配合 Lamb 优化器与 11e-3 高学习率seq 512 场景因上下文更长而把 batch 降为 32K、学习率降为 2e-3并启用了 rewarm。若要与预训练产出模型衔接下游任务验证请按 bert-finetuning.md 教程对 kernel 预训练模型做微调复现 SQuAD F1 指标。单卡吞吐结果Transformer Kernel 的收益可视化与上述大规模数据并行结果相辅相成的是单卡性能。下图文档原图存放于 docs/assets/images 目录展示了在 seq 128 与 seq 512 两种序列长度下经 DeepSpeed 优化的 BERT-Large 训练单卡吞吐对比教程记载相比 NVIDIA BERT 与 HuggingFace BERT 两个当时主流的 PyTorch 实现DeepSpeed 在 seq 128 与 seq 512 下分别达到约 64 与 53 teraflops 的单卡吞吐对应约 272 与 52 samples/second相对 NVIDIA BERT 提升最高约 28%、相对 HuggingFace BERT 提升最高约 62%同时支持在显存不溢出前提下把单卡 batch size 放大至约 1.8 倍。这些单卡能力的来源正是 Transformer Kernel 对多头注意力、LayerNorm、GELU 等算子的融合与算子级内存优化——相关配置字段与开关如pre_layer_norm、normalize_invertible、gelu_checkpoint、stochastic_mode的具体语义可在上一节及 transformer_kernel.md、transformer.py 中进一步研读。小结一套可以复用到任意 BERT 类项目的接入路径回看整个教程把 DeepSpeed 集成进 BERT 预训练代码库的完整动作可以浓缩为五步(1)用add_config_arguments()让train.py认识 DeepSpeed 参数(2)用deepspeed.initialize()把模型/优化器包成 DeepSpeed 引擎(3)把训练循环中的反向与参数更新换成引擎的backward()与step()(4)用save_checkpoint()/load_checkpoint()统一保存与恢复客户端状态和引擎内部状态(5)写一个声明批大小、优化器与 FP16 策略的 JSON 配置文件。再加上可选的 Transformer Kernel 融合层即可在同一份代码上获得从单卡吞吐优化到千卡级扩展的完整能力。这套方法与具体 BERT 实现NVIDIA/HuggingFace/Megatron-DeepSpeed解耦可平滑迁移到其他 Transformer 预训练工程。【免费下载链接】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），仅供参考