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 序列长度切分好的预训练样本目录，序列切分格式为 bnorick（bing_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 <json>（DeepSpeed 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(args=args, model=model.network, model_parameters=optimizer_parameters, dist_init_required=False) 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_required=False表示分布式初始化已由应用自行完成，引擎不再重复初始化；
• 配置既可以通过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.py）：

• backward(loss, retain_graph=False, scale_wrt_gas=True)：在非托管梯度累积模式下，backward()只做本地梯度累积，accumulation 边界处的梯度规约与优化器更新统一在step()内触发；同时会依据gradient_accumulation_steps()对 loss 进行缩放（loss = loss / gas），并针对 ZeRO、AMP 等路径做 loss scaling 处理；
• step(lr_kwargs=None)：执行参数更新与学习率调整，同时负责管理梯度清零与梯度累积边界逻辑；
• 引擎暴露的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.py：load_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 } }

配置要点逐条说明：

1. train_batch_size: 4096：期望的有效全局批大小（等效 batch size）为 4096；
2. train_micro_batch_size_per_gpu: 64：单卡显存可即时容纳的 micro batch size 为 64；DeepSpeed 会在引擎内部据此自动推导梯度累积步数gradient_accumulation_steps = train_batch_size / (micro_batch_size_per_gpu × data_parallel_world_size)，因此训练脚本里无需手工计算——这也正是引擎里gradient_accumulation_steps()属性存在的原因；
3. optimizer：使用 Adam 优化器，并给出学习率lr=2e-4、梯度裁剪阈值max_grad_norm=1.0、权重衰减weight_decay=0.01，bias_correction: false表示关闭 Adam 的偏置修正（典型用于 FP16 大 batch 训练）；
4. fp16：启用 FP16 混合精度训练。loss_scale: 0表示使用动态 loss scaling；initial_scale_power: 16表示动态 loss scaling 的初始缩放因子为2^16。

源码印证：配置中的这些键会被DeepSpeedConfig（deepspeed/runtime/config.py）解析，FP16 配置最终落到引擎的dynamic_loss_scale()/dynamic_loss_scale_args()相关逻辑（deepspeed/runtime/engine.py）；steps_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', default=False, action='store_true', help='Use 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_norm=True, attn_dropout_checkpoint=args.attention_dropout_checkpoint, normalize_invertible=args.normalize_invertible, gelu_checkpoint=args.gelu_checkpoint, stochastic_mode=True) 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。

四条重要使用注意事项（原样继承自教程）：

1. 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 或更小。
2. local_rank用于把 kernel 分配到正确设备。由于模型在此前已经执行过set_device()，因此这里不设置通常也没有问题。
3. stochastic_mode开启后性能更高但带有一定非确定性。预训练开启、微调关闭（微调需要可复现性，此点与 deepspeed/ops/transformer/transformer.py 中对stochastic_mode的建议一致）。
4. 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 张 V100（64 个 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）：

用于复现上述结果的训练配置摘要如下（详细脚本与配置位于配套 DeepSpeedExamples 仓库bing_bert目录）：

可复现性的两个关键点：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