DeepSpeed 实战：BingBertSQuAD 微调指南——将 SQuAD 问答微调无缝接入 DeepSpeed 并启用 Transformer Kernel 加速

本文是 DeepSpeed 官方教程「BingBertSQuAD Fine-tuning」的技术中文版，以仓库 docs/_tutorials/bert-finetuning.md 为主体骨架展开。教程演示如何为微软 BingBert 模型的 SQuAD v1.1 问答（Q&A）微调任务接入 DeepSpeed（该项目称为 BingBertSquad），内容包括最小化改动地改写训练脚本、通过 JSON 配置管理优化器与批大小、在 4 卡/批大小 24 的典型设定下复现 EM/F1 结果，以及在微调阶段开启 DeepSpeed Transformer Kernel 换取更高吞吐。读完本文，你可以独立完成 BERT 类模型到 DeepSpeed 训练范式（initialize / backward / step）的迁移，并能根据显存容量选择合适的 kernel 内存优化开关与预训练 checkpoint 格式。

1. BingBertSquad 教程定位与示例代码位置

BingBertSquad 是 DeepSpeed 官方示例集（DeepSpeedExamples）中的训练用例，教程中用到的关键脚本与目录位于 DeepSpeedExamples/training/BingBertSquad，相关入口如下：

• nvidia_run_squad_deepspeed.py —— 已经完成 DeepSpeed 改造的微调主脚本（复用 NVIDIA 官方 SQuAD 脚本结构）；
• nvidia_run_squad_baseline.py —— 未经 DeepSpeed 改造的基线版本，用于对比；
• run_squad_deepspeed.sh / run_squad_baseline.sh —— 两个一键启动脚本，分别驱动上面两个 Python 入口；
• deepspeed_bsz24_config.json —— 教程实验使用的 DeepSpeed JSON 配置（批大小 24）；
• utils.py —— 集中管理命令行参数定义与各类辅助函数；
• modeling.py 及 evaluate-v1.1.py —— BERT 建模与评测（EM/F1）脚本。

若需要在自己的环境中复现，可先克隆 DeepSpeed 发行源并初始化其示例子模块，随后进入示例目录：

git clone <DeepSpeed 仓库地址>
cd DeepSpeed
git submodule update --init --recursive
cd DeepSpeedExamples/training/BingBertSquad

本仓库聚焦于 DeepSpeed 核心引擎与算子实现，BingBertSquad 的建模与训练改造思路可以对照仓库内 deepspeed/ 包中的真实实现进行源码级验证，后文会在每个环节给出对应依据。

1.1 预备条件：数据与预训练 checkpoint

SQuAD v1.1 数据：需要下载两个官方数据文件并放入数据目录：

• 训练集 train-v1.1.json
• 验证集 dev-v1.1.json

从 SQuAD 官方站点获取上述两个文件即可。它们将分别用于微调训练与评测。

预训练 BERT checkpoint：三种来源任选其一：

1. DeepSpeed 预训练产出：本教程使用 DeepSpeed 官方 BERT 预训练教程（bert-pretraining.md）中产出的 checkpoint 160；
2. HuggingFace：bert-large-uncased-whole-word-masking（PyTorch 权重 *_pytorch_model.bin 与配套 JSON config）；
3. TensorFlow/Google：Bert-large-uncased-L-24_H-1024_A-16（Google BERT 预训练权重 zip，解压后为 bert_model.ckpt.* 系列文件）。

需要注意 checkpoint 与当前脚本使用的词表（vocabulary）必须一致，BingBertSquad 启动时会做一致性校验。

2. 运行 BingBertSquad：DeepSpeed 版与基线版

教程提供两个 shell 启动脚本，参数完全一致，便于对照：

启动脚本	驱动脚本	说明
run_squad_deepspeed.sh	nvidia_run_squad_deepspeed.py	DeepSpeed 启用版，本教程主角
run_squad_baseline.sh	nvidia_run_squad_baseline.py	未启用 DeepSpeed 的基线

统一接收 4 个位置参数：

bash run_squad_deepspeed.sh <NUM_GPUS> <PATH_TO_CHECKPOINT> <PATH_TO_DATA_DIR> <PATH_TO_OUTPUT_DIR>

各参数含义：

1. <NUM_GPUS>：参与训练的 GPU 数量；
2. <PATH_TO_CHECKPOINT>：第 1.1 节准备的预训练 checkpoint 路径；
3. <PATH_TO_DATA_DIR>：存放 train-v1.1.json 与 dev-v1.1.json 的数据目录；
4. <PATH_TO_OUTPUT_DIR>：训练产物（模型与 predictions.json）输出目录。

3. DeepSpeed 集成：一个普通微调脚本如何被改造

nvidia_run_squad_deepspeed.py 已经完成全部改造，读者无需自行修改即可运行。理解其改造点，对把任意 PyTorch 训练脚本迁移到 DeepSpeed 具有直接的参考价值，整个改造只涉及四个层面：配置注入、参数解析、初始化与训练循环替换。

3.1 配置：deepspeed_bsz24_config.json

DeepSpeed 采用 JSON 配置驱动运行时行为。训练时除显式添加 --deepspeed 开关外，还需用 --deepspeed_config deepspeed_bsz24_config.json 指定配置文件。教程实验（Table 1）采用的微调配置如下，这是后续所有复现实验的基准：

参数	值
Total batch size（总批大小）	24
Train micro batch size per GPU（每卡微批）	3
Optimizer	Adam
Learning rate	3e-5
Sequence-length（序列长度）	384
Weight-decay	0.0
Epoch count	2

以 DeepSpeed 配置标准键名表达，一个与 Table 1 吻合的可参考 JSON 形态为：

{
  "train_batch_size": 24,
  "train_micro_batch_size_per_gpu": 3,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 3e-5,
      "weight_decay": 0.0
    }
  },
  "fp16": {
    "enabled": true
  }
}

从配置语义看，DeepSpeed 会依据 GPU 数量、train_batch_size 与 train_micro_batch_size_per_gpu 的关系推导梯度累积步数：本教程在 4 张 GPU 上运行时，24 / (3 × 4) = 2，即每累积 2 个微批的梯度做一次参数更新。也就是说用户无需在代码里手写梯度累积循环，这正是后续 model.backward/model.step 抽象能成立的原因。

3.2 参数解析：deepspeed.add_config_arguments()

改造第一步是在主入口 main() 开头，用 DeepSpeed 提供的辅助函数向原有参数解析器注入 --deepspeed、--deepspeed_config 等命令行参数：

parser = get_argument_parser()
# Include DeepSpeed configuration arguments
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()

其中 get_argument_parser() 来自示例的 utils.py，模型相关的全部选项及其描述也在该文件中维护。仓库侧的实现位于 deepspeed/init.py：add_config_arguments() 内部调用 _add_core_arguments()，后者向解析器注册一个名为 “DeepSpeed” 的参数组，核心即：

• --deepspeed（布尔开关，default=False，用于开启 DeepSpeed，仅作为用户侧辅助标志）；
• --deepspeed_config（字符串，指向 DeepSpeed JSON 配置文件的路径，见 deepspeed/init.py）。

--deepspeed_config 随后会被读取并加载为 DeepSpeedConfig，因此在命令行参数层，用户几乎不需要关心 JSON 里每一项的底层含义。

3.3 训练循环的四段式改造

(1) 初始化（Initialization）

DeepSpeed 提供 initialize() 一站式封装 model、optimizer、LR scheduler 与 data loader。BingBertSquad 中只需传入模型与待优化参数分组，即可拿到已被 DeepSpeed 引擎包裹的 model（返回的其实是 DeepSpeedEngine）与自动构建的优化器：

model, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=optimizer_grouped_parameters
)

仓库侧，initialize() 的完整签名与返回约定见 deepspeed/init.py：它接受 args、model、optimizer、model_parameters、training_data、lr_scheduler、config 等参数，返回 engine, optimizer, training_dataloader, lr_scheduler 四元组。函数内部会完成分布式通信初始化（dist.init_distributed），并将 args.deepspeed_config 解析为配置对象（deepspeed/init.py）。注意：传入 --deepspeed_config 与通过 config= 参数显式传配置二者只能取其一，同时给出会触发断言。

(2) 前向传播（Forward pass）

前向在基线与 DeepSpeed 版本中完全一致：

loss = model(input_ids, segment_ids, input_mask, start_positions, end_positions)

SQuAD 问答任务的损失由答案起始/结束位置两个 head 的目标联合给出。

(3) 反向传播（Backward pass）

这是改写幅度最大的地方：

• 基线版本：在梯度累积边界必须显式处理跨卡 all-reduce。FP16 下要做 enable_need_reduction() 之后再 optimizer.backward(loss)，并自行维护动态 loss scaling；FP32 下调用 loss.backward()。
• DeepSpeed 版本：只需一行 model.backward(loss)。

仓库引擎中，DeepSpeedEngine.backward() 的实现位于 deepspeed/runtime/engine.py，它会在引擎内部统一完成梯度累积、跨 GPU all-reduce（按梯度累积步数 scale_wrt_gas 缩放）与混合精度 loss scaling 的管理，把原先散落在用户代码里的分布式细节全部收敛掉。

(4) 参数更新（Weight updates）

• 基线版本：FP16 需显式选择 FusedAdam 并处理动态 loss scaling，FP32 使用 BertAdam；更新时刻还要手动执行 optimizer.step() 与 optimizer.zero_grad()。
• DeepSpeed 版本：优化器在 initialize() 时已按 JSON 配置构建好，因此训练循环内只需：

model.step()

引擎的 DeepSpeedEngine.step() 见 deepspeed/runtime/engine.py，负责在合适的梯度累积边界触发优化器 step 并归零梯度（scheduler 的 step 也在引擎内部统一驱动）。

完成以上四步，一个模型的 DeepSpeed 移植即告完成——BingBertSquad 的改造完全遵循这一模式，未改动任何模型结构。

3.4 评测：计算 EM 与 F1

训练结束后，用官方评测脚本对输出目录中的预测结果打分：

python evaluate-v1.1.py <PATH_TO_DATA_DIR>/dev-v1.1.json <PATH_TO_DATA_DIR>/predictions.json

predictions.json 由微调脚本在 <PATH_TO_OUTPUT_DIR> 下生成；评测以 SQuAD 标准的 Exact Match（EM） 与 F1 为指标。

3.5 教程实验的微调结果

教程在 DGX-2 节点上用 4 张 GPU、总批大小 24 训练 2 个 epoch 得到以下结果（所有实验学习率均为 3e-5，并通过网格挑选了最优随机种子：HuggingFace 与 TensorFlow 模型分别使用 seed 9041 与 19068）：

预训练来源	模型	精度	EM	F1
TensorFlow	Bert-large-uncased-L-24_H-1024_A-16	FP16	84.13	91.03
HuggingFace	Bert-large-uncased-whole-word-masking	FP16	87.27	93.33

说明：上述数值为官方教程记录的实验数据；使用 DeepSpeed 预训练 checkpoint 并配合后文 kernel 微调可获得的 F1 见第 5.4 节。读者实际运行结果会因软硬件、随机种子与数据预处理细节而略有差异。

4. 启用 DeepSpeed Transformer Kernel 提升吞吐

Transformer Kernel 是 DeepSpeed 针对 Transformer 层打造的高性能融合算子：它在单个算子内融合多头自注意力与前馈网络的关键计算，减少 kernel launch 与显存中间张量，从而在单卡吞吐与多卡扩展性上同时获益（机制细节可参阅 transformer_kernel.md）。在微调阶段，该 kernel 不仅可以用于 DeepSpeed 自身预训练的模型，也可作用于 TensorFlow 与 HuggingFace 产出的 checkpoint。

4.1 开关参数与层替换

utils.py 中已经预置了开关参数：

parser.add_argument(
    '--deepspeed_transformer_kernel',
    default=False,
    action='store_true',
    help='Use DeepSpeed transformer kernel to accelerate.'
)

在建模源码的 BertEncoder 类中，当 --deepspeed_transformer_kernel 开启时，会用 DeepSpeedTransformerLayer 构造整个编码器（对每一层 deepcopy 一份独立实例并赋予自增的 layer_id），否则退回原生 BertLayer：

if args.deepspeed_transformer_kernel:
    from deepspeed import DeepSpeedTransformerLayer, \
        DeepSpeedTransformerConfig, DeepSpeedConfig

    ds_config = DeepSpeedConfig(args.deepspeed_config)

    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,
        seed=args.seed,
        fp16=ds_config.fp16_enabled
    )
    self.layer = nn.ModuleList([
        copy.deepcopy(DeepSpeedTransformerLayer(i, cuda_config))
        for i in range(config.num_hidden_layers)
    ])
else:
    layer = BertLayer(config)
    self.layer = nn.ModuleList([
        copy.deepcopy(layer)
        for _ in range(config.num_hidden_layers)
    ])

因此模型源码必须能拿到 args，所有配置均来自 DeepSpeed 配置文件与命令行参数。DeepSpeedTransformerLayer 要求的是独立的 layer 实例（每层需要不同的 layer_id），这与源码中静态计数器 DeepSpeedTransformerLayer.layer_id 自增的设计一致，见 deepspeed/ops/transformer/transformer.py。

关于 batch_size 的关键约束：batch_size 是输入数据的最大批大小上限。所有微调训练数据或预测数据的批大小都不能超过该阈值，否则会抛异常。在 DeepSpeed 配置文件中，该值由 train_micro_batch_size_per_gpu 给出；例如 train_micro_batch_size_per_gpu 设为 8，则预测时的 --predict_batch_size 也必须设为 8，二者需要保持同步。

4.2 TransformerConfig 关键字段与 kernel 内存优化开关

上述 DeepSpeedTransformerConfig 构造时的 fp16、seed 等字段，以及 kernel 内部使用的 pre_layer_norm、normalize_invertible、gelu_checkpoint、attn_dropout_checkpoint、stochastic_mode 等参数，其完整定义可对照 deepspeed/ops/transformer/transformer.py。与本节性能调优直接相关的三个内存优化开关语义如下（默认均为 False）：

开关	作用
normalize_invertible	使用可逆 LayerNorm 执行，释放输入激活的内存
gelu_checkpoint	对 GELU 激活输出做重计算式 checkpoint 以省显存
attn_dropout_checkpoint	对注意力 dropout 做重计算式 checkpoint 以省显存

开启开关后，kernel 在 CUDA 层创建实例时会把对应布尔量透传给底层 kernel（见 deepspeed/ops/transformer/transformer.py）。直观地讲，这些开关牺牲少量计算换取显存容量，从而允许更大的微批。

4.3 加载 HuggingFace 与 TensorFlow 预训练模型

BingBertSquad 同时支持三种 checkpoint 格式，示例目录结构如下：

[test/huggingface]
bert-large-uncased-whole-word-masking-config.json
bert-large-uncased-whole-word-masking-pytorch_model.bin

[test/tensorflow]
bert_config.json
bert_model.ckpt.data-00000-of-00001
bert_model.ckpt.index
bert_model.ckpt.meta

加载由三个参数控制：

1. --model_file：指向预训练模型权重文件；
2. --ckpt_type：checkpoint 类型，TF 表示 TensorFlow，HF 表示 HuggingFace，默认值 DS 表示 DeepSpeed；
3. --origin_bert_config_file：BERT 结构配置文件，通常与 model_file 同目录存放。

在 run_squad_deepspeed.sh 中分别追加以下片段即可运行对应示例：

[HuggingFace]

--model_file test/huggingface/bert-large-uncased-whole-word-masking-pytorch_model.bin \
--ckpt_type HF \
--origin_bert_config_file test/huggingface/bert-large-uncased-whole-word-masking-config.json \

[TensorFlow]

--model_file /test/tensorflow/bert_model.ckpt \
--ckpt_type TF \
--origin_bert_config_file /test/tensorflow/bert_config.json \

使用外部预训练模型时有三个注意事项：

1. 使用 HuggingFace 或 TensorFlow 预训练模型，必须开启 --deepspeed_transformer_kernel；
2. 它们属于 Post-Layer-Norm 结构，因此不能与 --preln（Pre-Layer-Norm）标志混用；
3. BingBertSquad 会校验预训练模型的词表大小，出现任何不匹配都无法运行，建议使用上文描述风格的官方 checkpoint 或 DeepSpeed 产出的 bing_bert checkpoint。

5. 性能调优：面向不同显存的微批与内存开关选择

在总批大小固定为 24（Table 1）的前提下，可以通过增大每卡微批大小来提高吞吐。官方教程在 NVIDIA V100（16GB 与 32GB）上做了系统性的微批扫描实验（使用 DeepSpeed Transformer Kernel），结论是：微批越大、单卡吞吐越高。

V100 16GB 上的 samples/second（Table 2）：

每卡微批	PyTorch	DeepSpeed	加速比（x）
4	36.34	50.76	1.4
6	OOM	54.28	1.5
8	OOM	54.16	1.5

16GB 场景下，相比 PyTorch 最多可取得约 1.5x 加速，同时单卡支持 2 倍于 PyTorch 的批大小。

V100 32GB 上的 samples/second（Table 3）：

每卡微批	PyTorch	DeepSpeed	加速比（x）
4	37.78	50.82	1.3
6	43.81	55.97	1.3
12	49.32	61.41	1.2
24	OOM	60.70	1.2
32	OOM	63.01	1.3

32GB 场景下单卡微批可做到 32，约为 PyTorch 容量的 2.6 倍，同时端到端微调仍保持约 1.3x 加速。注意：对 PyTorch 出现 OOM（内存溢出）的场景，加速比使用 DeepSpeed 可支持场景下的最佳 samples/second 计算。

达到不同微批所需的内存优化开关（Table 4）——参考“DeepSpeed Transformer Kernel”教程设置：

每卡微批	NVIDIA V100（32GB）	NVIDIA V100（16GB）
> 4	-	normalize_invertible
> 6	-	attn_dropout_checkpoint, gelu_checkpoint
> 12	normalize_invertible, attn_dropout_checkpoint	OOM
> 24	gelu_checkpoint	OOM

实操建议：微批从 3 起步逐步上调（如 24 甚至更高）前，先根据上表在对应显存容量下补齐所需内存开关；16GB 显卡上微批上限约为 6–8，超过后即使开启全部开关也仍会 OOM，需要切换到 32GB 显存或降低总批大小。

6. 微调 DeepSpeed Transformer Kernel 预训练模型：dropout 设定

对“使用 DeepSpeed Transformer Kernel 与 bert-pretraining.md 中的 Fast-BERT 预训练配方”产出的模型做微调，官方教程记录 F1 约为 90.5，并预期若预训练时间超过教程建议的轮数，F1 还会进一步提升。

这一结果需要针对 dropout 做特殊设定，以适配预训练使用的 kernel 模式（确定性 deterministic 与随机性 stochastic 两种模式的取舍详见 transformer_kernel.md）：

• 微调阶段只使用确定性（deterministic） transformer，以保证微调结果可复现（源码注释也明确建议：下游微调任务应关闭 stochastic 模式以获得可复现结果，见 deepspeed/ops/transformer/transformer.py）；
• dropout 取值取决于预训练模式：预训练使用确定性 transformer 时，微调沿用预训练的 dropout 比例 0.1；预训练使用随机性（stochastic）transformer 时，由于微调缺少随机噪声，需略微调高 dropout 加以补偿。

预训练模式	Dropout ratio
Deterministic（确定性）	0.1
Stochastic（随机性）	0.12 - 0.14

7. 小结：一条可复用的“传统模型 + DeepSpeed”迁移路径

从 BingBertSquad 案例中可以提炼出一条对任意 PyTorch 训练脚本都适用的 DeepSpeed 迁移路径：

1. 配置先行：编写 JSON（train_batch_size、train_micro_batch_size_per_gpu、optimizer/lr、fp16 等），启动时用 --deepspeed_config 注入；
2. 解析扩展：调用 deepspeed.add_config_arguments(parser) 获得 --deepspeed/--deepspeed_config 参数；
3. 初始化替换：用 deepspeed.initialize(args=args, model=model, model_parameters=...) 一次性完成模型包裹与优化器构建（deepspeed/init.py）；
4. 训练循环收敛：前向不变，用 model.backward(loss) 与 model.step() 取代手写的梯度累积、all-reduce、动态 loss scaling 与 optimizer.step()/zero_grad()（deepspeed/runtime/engine.py、deepspeed/runtime/engine.py）；
5. 吞吐加速：需要更高吞吐时启用 --deepspeed_transformer_kernel 替换 Transformer 层（deepspeed/ops/transformer/transformer.py），并按显存容量与目标微批配置内存优化开关。

该路径不依赖任何特殊模型结构，正是 DeepSpeed 被设计为“以极小代码改动获得分布式训练与显存优化收益”的典型体现。SQuAD 问答微调只是第一步，同一套范式同样适用于后续教程涉及的分类、生成等各类 Transformer 下游任务。