DeepSpeed Transformer Kernels API 指南：使用 DeepSpeedTransformerConfig 与 DeepSpeedTransformerLayer 构建高效 BERT Transformer 层

Transformer 层是现代自然语言处理等序列模型中通用的基础组件，其计算效率直接决定预训练与微调任务的吞吐。DeepSpeed 通过自定义 CUDA kernel 将 Transformer 层的多算子融合执行，并把这一能力封装为两个公开 Python API：DeepSpeedTransformerConfig（配置类）与 DeepSpeedTransformerLayer（层模块类），帮助用户构建面向预训练与微调的高效 BERT Transformer 层。本文以官方 API 文档 kernel.rst 为主体，结合仓库内源码、示例与测试，完整讲解配置参数语义、层初始化机制、底层前反向执行流程与内存优化手段，读完即可在自有 BERT 类模型中启用 DeepSpeed Transformer kernel 并正确调参。

Transformer Kernel API 概览与定位

根据 kernel.rst，Transformer kernel API 用于创建 BERT transformer 层以实现更高效的预训练（pre-training）与微调（fine-tuning），其组成部分为：

• Transformer 层配置：DeepSpeedTransformerConfig；
• Transformer 层模块初始化：DeepSpeedTransformerLayer。

两个类均在模块顶层 deepspeed 命名空间直接导出，deepspeed/init.py 中的导入语句为：

from .ops.transformer import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig

因此使用时的标准导入方式为 from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig，无需关心底层 kernel 的编译细节。两个类的完整实现位于 deepspeed/ops/transformer/transformer.py：

符号	基类/类型	行号	职责
TransformerConfig	普通类	L19-L31	保存层几何/超参基础字段
DeepSpeedTransformerConfig	TransformerConfig 子类	L34-L140	完整 kernel 配置，含内存/精度开关
DeepSpeedTransformerFunction	torch.autograd.Function	L143-L293	kernel 前向/反向与自动求导桥接
DeepSpeedTransformerLayer	nn.Module	L296-L414	权重定义、kernel 层创建与前向入口

从源码结构可以推断，整个 API 的执行链路为：DeepSpeedTransformerLayer.forward() → DeepSpeedTransformerFunction.apply() → C++ 编译出的 kernel 模块 transformer_cuda_module/stochastic_transformer_cuda_module（按是否启用 stochastic_mode 选择，由 TransformerBuilder/StochasticTransformerBuilder 加载）。CUDA 算子构建逻辑见 deepspeed/ops/op_builder/transformer.py，对应 C++/CUDA 实现位于 csrc/transformer/。

注意：kernel 接口的官方使用场景文档为 BERT 预训练教程。kernel.rst 明确指引读者参阅 BERT pre-training tutorial 获取完整用法，本文后续章节会给出基于当前仓库源码的准确集成示例。

DeepSpeedTransformerConfig：完整参数语义与默认值

DeepSpeedTransformerConfig 是驱动整个 kernel 行为的配置对象，继承自 TransformerConfig 并持有运行时开关。transformer.py L89-L127 给出了构造函数的完整签名与默认值，下文按官方文档分类整理参数语义。

通用结构参数（层的几何配置）

这类参数描述 Transformer 层的基本形状，由用户依据模型架构填写：

参数	默认值	含义
batch_size	-1	每块 GPU 上运行 kernel 的最大 batch size
hidden_size	-1	Transformer 层隐层维度
intermediate_size	-1	前馈网络中间层维度；若传入值 ≤ 0，则自动取 4 * hidden_size（见 L110-L113）
heads	-1	自注意力头数
attn_dropout_ratio	-1	注意力输出的 dropout 比例
hidden_dropout_ratio	-1	Transformer 输出的 dropout 比例
num_hidden_layers	-1	模型中 Transformer 层总数（影响权重初始化缩放与 layer_id 编号范围）
initializer_range	-1	参数初始化标准差（BERT 约定值，如 0.02）

架构与精度参数

参数	默认值	含义
layer_norm_eps	1e-12	LayerNorm 数值稳定项 epsilon
fp16	False	是否启用半精度计算；控制调用 forward_fp16/backward_fp16 还是 forward_fp32/backward_fp32 kernel
pre_layer_norm	True	选择 Pre-LN 还是 Post-LN 架构。官方 BERT 预训练示例推荐 True
adjust_init_range	True	是否在初始化时对残差路径权重缩小标准差：output_std = initializer_range / sqrt(2.0 * num_layers)（实现见 L379-L397）
training	True	训练（而非推理）模式；DeepSpeedTransformerLayer.forward() 会在每次前向时将其同步为 self.training

运行环境参数

参数	默认值	含义
local_rank	-1	运行 kernel 的 GPU 编号；当模型未自行调用 set_device 时需要设置，kernel 层据此把工作放到正确的设备上（见 L320-L321）。如果模型此前已设置当前设备，可不设置
seed	-1	dropout 层的随机种子
return_tuple	False	为 True 时前向返回 (output,) 元组形式（便于 Hugging Face 风格模型使用 outputs[0]）

高性能开关：stochastic_mode

stochastic_mode 默认 False。启用后，层的初始化与前反向会改用 StochasticTransformerBuilder 编译的随机性 kernel（见 L364-L377），运行更快，但存在一定程度的非确定性，多次运行结果可能不同。官方实践建议：预训练任务（如 BERT）可开启以获得更高吞吐且精度基本不受影响；下游微调任务为保证可复现，建议关闭、走常规 kernel。需要明确的是，kernel.rst 与源码仅陈述该模式的非确定性与适用建议，未给出具体加速百分比数据。

内存优化开关

参数	默认值	含义
normalize_invertible	False	启用可逆 LayerNorm 执行，丢弃 LayerNorm 的输入激活（反向时仅用输出激活即可恢复梯度）
gelu_checkpoint	False	对 GELU 激活输出做 checkpointing，丢弃其输入激活以省显存；需要更大 batch 时可开启
attn_dropout_checkpoint	False	对注意力 dropout 做 checkpointing，丢弃其输入激活以省显存

这三个开关对激活显存的影响与反向时保存张量的取舍，将在「内存优化标志位的源码级原理」一节结合 DeepSpeedTransformerFunction.backward 详述。

便捷构造方法

除直接调用构造函数外，配置类还提供了两个反序列化方法（L129-L140）：

config = DeepSpeedTransformerConfig()
config = DeepSpeedTransformerConfig.from_dict(json_object)   # 从 dict 填充
config = DeepSpeedTransformerConfig.from_json_file(json_file)  # 从 UTF-16 JSON 文件读取

其中 from_dict 遍历 dict 键值写入 config.__dict__；from_json_file 以 UTF-16 编码读取 JSON 文件后调用 from_dict。单元测试（如 test_accelerator_forward.py）采用的即是先 DeepSpeedTransformerConfig() 再逐字段赋值的写法。

DeepSpeedTransformerLayer：层实例化、权重初始化与 kernel 层创建

DeepSpeedTransformerLayer 继承 torch.nn.Module。根据 docstring 与构造实现 L296-L377，其关键设计如下。

静态 layer_id 自动编号

类维护一个静态计数器 layer_id（初始为 0）。每次实例化一个层，都会把当前计数赋给 config.layer_id 并将计数器加一：

self.config.layer_id = DeepSpeedTransformerLayer.layer_id
DeepSpeedTransformerLayer.layer_id += 1

因此，若一个 BERT 模型有 24 个 Transformer 层，layer_id 会从 0 递增到 23，与 kernel 内部按层索引管理的工作区一一对应。

权重定义与初始化

当未传入 initial_weights/initial_biases（这两个参数官方注释为“仅单元测试使用”）时，层会自动创建全套参数（L323-L336）：

参数	形状	说明
attn_qkvw	(3*hidden_size, hidden_size)	融合的 Q/K/V 投影权重
attn_qkvb	(3*hidden_size,)	融合的 Q/K/V 偏置（初始化为 0）
attn_ow / attn_ob	(hidden_size, hidden_size) / (hidden_size,)	注意力输出投影
attn_nw / attn_nb	(hidden_size,)	注意力后 LayerNorm 参数（gamma 初始化为 1，beta 为 0）
inter_w / inter_b	(intermediate_size, hidden_size) / (intermediate_size,)	前馈第一层
output_w / output_b	(hidden_size, intermediate_size) / (hidden_size,)	前馈第二层
norm_w / norm_b	(hidden_size,)	输出 LayerNorm 参数

初始化逻辑（init_transformer_weights，L379-L397）：普通权重用均值 0、标准差 initializer_range 的正态分布采样；残差路径上的输出权重（attn_ow、output_w）在 adjust_init_range=True 时改用 initializer_range / sqrt(2.0 * num_layers) 的标准差，这是官方在注释中所称“考虑残差路径上的累积”的初始化调整。

kernel 模块加载与层创建

第一次构造层时，模块按需加载 CUDA 实现（模块级缓存 transformer_cuda_module）：

if transformer_cuda_module is None and not self.config.stochastic_mode:
    transformer_cuda_module = TransformerBuilder().load()
if stochastic_transformer_cuda_module is None and self.config.stochastic_mode:
    stochastic_transformer_cuda_module = StochasticTransformerBuilder().load()

随后调用 kernel 侧 API 在 GPU 上创建对应层（fp16 走 create_transformer_layer_fp16，否则 create_transformer_layer_fp32），并把 layer_id、batch_size、hidden_size、heads、intermediate_size、dropout 比例、layer_norm_eps、seed、pre_layer_norm、内存优化开关与 stochastic_mode 一并下传（L370-L377）。

前向与反向执行：autograd Function 桥接 kernel 的细节

DeepSpeedTransformerLayer.forward()（L399-L414）会将 torch.is_grad_enabled() 与 self.training 同步进配置，再一次性把 hidden_states、attention_mask 及全部参数交给 DeepSpeedTransformerFunction.apply。DeepSpeedTransformerFunction 是基于 torch.autograd.Function 的薄封装，实际计算全部发生在 C++ kernel。从源码可以提炼以下工程细节：

1. fp16/fp32 与随机 kernel 分派（L149-L150）：前向按 stochastic_mode 先选模块，再按 fp16 选 forward_fp16 或 forward_fp32；反向对应 backward_fp16/backward_fp32（L252-L253）。

2. 序列长度对齐 16（L152-L160）：若序列维不是 16 的整数倍，输入会被随机张量填充（掩码以 -10000 填充以屏蔽注意力），计算结束后用 torch.narrow 裁回原长度（L227-L228）；反向阶段对 grad_output 做同样的补齐与裁剪（L239-L241、L288-L289）。

3. 反向中间张量取舍与配置开关严格对应（L192-L225）：训练且梯度启用时，ctx 依据各开关决定保存哪些中间结果——例如 attn_dropout_checkpoint=True 时不保存 ctx_bufB，gelu_checkpoint=True 时不保存 gelu_inp，normalize_invertible=True 且 Pre-LN 时不再保存原始 input。这正是三个内存优化开关能省显存的直接原因。

4. 上下文显存主动释放（L269-L286）：反向完成后把 ctx 中所有中间张量置 None，作者注释称之为“一种有效释放上下文显存的方式”。

5. 梯度 Hook（仅测试）（L171-L190）：当 forward 传入 grads 列表时注册 register_hook 收集 Q/K/V 与各层权重的梯度，用于与参考实现做逐参数对比，是单元测试专用路径。

在 BERT 预训练/微调中集成 Transformer Kernel

步骤一：构造配置并实例化层

kernel.rst 把使用细节指向 BERT pre-training tutorial。结合该教程的意图与当前仓库 构造函数签名，可得到一段可直接运行的准确集成示例（所有配置取自 DeepSpeed 配置文件与命令行参数）：

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,
    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,                 # 官方示例使用 Pre-LN BERT-Large
    attn_dropout_checkpoint=args.attention_dropout_checkpoint,
    normalize_invertible=args.normalize_invertible,
    gelu_checkpoint=args.gelu_checkpoint,
    stochastic_mode=True,                # 预训练开启；微调建议关闭以保证可复现
)

layer = DeepSpeedTransformerLayer(cuda_config)
# 一个编码器由 num_hidden_layers 层拷贝组成，layer_id 自动 0..num_hidden_layers-1
self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])

实操要点（原教程 Note，均保留自 docs/_tutorials/bert-pretraining.md）：

• batch_size 是上限而非精确值：它表示输入数据的最大 batch，微调/预测的数据 batch 不应超过该阈值，否则抛出异常。DeepSpeed 配置文件中对应 train_micro_batch_size_per_gpu。
• local_rank 只需在模型未主动 set_device 时设置：若模型此前已设置当前设备，此处可不传（默认 -1）。
• stochastic_mode 预训练开、微调关：预训练开启可获得更高吞吐且不影响收敛精度；微调阶段关闭以保证结果可复现。
• kernel 权重是独立参数，checkpoint 强耦合 kernel：用 Transformer kernel 训练产出的 checkpoint 文件必须由启用 kernel 的模型加载（如后续微调时同样开启 kernel），否则参数结构不一致。

步骤二：通过命令行开关启用 kernel

在训练脚本中为内核能力预留开关（默认关闭以便对照回退）：

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

启动时（沿用 DeepSpeed launcher 与标准 deepspeed 配置项）：

deepspeed deepspeed_train.py \
    --cf bert_large_lamb.json \
    --max_seq_length 512 \
    --deepspeed \
    --deepspeed_transformer_kernel \
    --deepspeed_config deepspeed_bsz32K_lamb_config_seq512.json \
    --attention_dropout_checkpoint \
    --max_steps 32 \
    --print_steps 100

兼容性提醒

从 docs/_tutorials/transformer_kernel.md 可知：当前 DeepSpeed Transformer Kernel 不支持 Sparse Attention，若模型使用 Sparse Attention 需要关闭 Transformer kernel。这是官方文档明确声明的适用边界。

内存优化标志位的原理与选型参考

「高性能大 batch 训练」需要在激活显存上精打细算。DeepSpeed 在 kernel 内部实现了三类省显存手段，均通过配置开关暴露（详见 docs/_tutorials/transformer_kernel.md）：

• normalize_invertible（可逆 LayerNorm）：丢弃 LayerNorm 输入激活。由于 kernel 实现了“仅用输出激活即可反推参数与输入梯度”的优化，该输入无需保留。反向代码中体现为：Pre-LN 且可逆时不再 save_for_backward 原始 input，而是保存归一化中间量 inp_norm 并以此为输入计算（L193-L202）。
• attn_dropout_checkpoint（注意力 dropout 激活重计算）：丢弃注意力 dropout 的输入，反向用 dropout mask 恢复；对应反向时以 soft_inp 代替未保存的 ctx_bufB（L206-L207、L260）。
• gelu_checkpoint（GELU 激活重计算）：丢弃 GELU 输入，反向时以 ff2_inp 代替未保存的 gelu_inp（L217-L218、L262）。

官方教程给出的经验法则是：后两者虽然引入重计算开销，但实测对整体性能的影响“可忽略”，而通过省下的显存换取更大 batch 后，端到端训练效率反而上升。对于何时开启哪些开关，教程基于 BERT-Large + V100 32GB 给出如下选型表（从 docs/_tutorials/transformer_kernel.md 完整继承）：

Micro-batch size	128 sequence-length	512 sequence-length
> 12	-	attn_dropout_checkpoint
> 16	-	normalize_invertible、gelu_checkpoint
> 80	normalize_invertible	OOM
> 112	attn_dropout_checkpoint	OOM
> 128	gelu_checkpoint	OOM

需要说明：该选型表是仓库文档基于特定硬件（V100 32GB）与 BERT-Large 配置给出的指导性经验，具体阈值应结合自身 GPU 显存与模型规模重新实测。

源码与测试佐证：如何验证你的 kernel 集成

仓库在 tests/unit/ops/accelerators/ 下提供了 kernel 的等价性验证用例，可作为复现与回归参照：

• test_accelerator_forward.py：构建参考 BERT 编码器（Pre-LN 来自 unit/modelingpreln，Post-LN 来自 unit/modeling）与 DeepSpeed kernel 编码器，用同一组初始权重对比前向输出。覆盖 fp16/fp32、Pre-LN/Post-LN、数百种 (batch_size, hidden_size, seq_len, heads, num_layers) 组合，甚至包括小 batch（构造 batch 8 实际跑 batch 3/7）与非 16 对齐序列长（21、51、53、119、381、509 等）；stochastic_mode 用例放宽误差到 7e-2，印证其非确定性与精度特性。
• test_accelerator_backward.py：对应反向传播与逐参数梯度的等价性校验，验证了前文所述 autograd Function 的反向链路。

若要在自己的环境复现，测试均通过 TransformerBuilder 是否兼容决定是否跳过（deepspeed.ops.__compatible_ops__），且依赖 deepspeed 顶层导出，即测试先决条件与 kernel API 的正式使用方式一致。

小结

DeepSpeedTransformerConfig + DeepSpeedTransformerLayer 构成了 DeepSpeed Transformer kernel 的完整 Python 面：前者以 20 余项参数覆盖层几何、精度（fp16/fp32）、架构（Pre-LN/Post-LN）、随机 kernel、可逆 LayerNorm、激活 checkpointing 等全部可调面；后者负责 layer_id 编号、参数自动初始化、按配置加载并创建 GPU kernel 层，并通过 torch.autograd.Function 与 C++ 前反向无缝衔接。集成该 API 只需两步——在模型编码器中用 kernel 层替换普通 BERT 层、在启动命令中加入开关——并遵循「batch 上限」「checkpoint 与 kernel 强耦合」「预训练开 / 微调关 stochastic_mode」等约束即可在 BERT 类模型上获得高效训练。如需深入了解数值实现，可继续阅读 deepspeed/ops/transformer/transformer.py、csrc/transformer/ 及 deepspeed/ops/op_builder/transformer.py 对应源码。