DeepSpeed BERT 预训练实战：把 DeepSpeed 无缝接入现有 Transformer 预训练代码库

2026-09-08 11:22:04作者：蔡丛锟

本篇技术指南以 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
  }
}

配置要点逐条说明：

train_batch_size: 4096：期望的有效全局批大小（等效 batch size）为 4096；
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() 属性存在的原因；
optimizer：使用 Adam 优化器，并给出学习率 lr=2e-4、梯度裁剪阈值 max_grad_norm=1.0、权重衰减 weight_decay=0.01，bias_correction: false 表示关闭 Adam 的偏置修正（典型用于 FP16 大 batch 训练）；
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。

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

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 张 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）：

节点数	V100 GPU 数	训练耗时
1 DGX-2	16	33 hr 13 min
4 DGX-2	64	8 hr 41 min
16 DGX-2	256	144 min
64 DGX-2	1024	44 min

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

参数	128 序列长度	512 序列长度
总 batch size	64K	32K
单卡 micro batch size	64	8
优化器	Lamb	Lamb
学习率	11e-3	2e-3
初始学习率（`lr_offset`）	10e-4	0.0
Min Lamb 系数	0.01	0.01
Max Lamb 系数	0.3	0.3
LR 调度器	`warmup_exp_decay_exp`	`warmup_exp_decay_exp`
Warmup 比例	0.02	0.02
Decay 率	0.90	0.90
Decay 步数	250	150
最大训练步数	7500	7500
Rewarm 学习率	N/A	True
输出检查点数	150	160-162
样本数	403M	18-22M
Epoch 数	150	160-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 预训练工程。