DeepSpeed Transformer Kernel 使用指南：配置、内存优化与启动实战

Transformer 层几乎出现在所有现代序列建模模型中，其训练效率直接决定了 NLP 预训练与微调的吞吐量。DeepSpeed 为此提供了专门定制的 Transformer Kernel（深度融合算子），本教程基于 docs/_tutorials/transformer_kernel.md，讲解如何在训练脚本中实例化 DeepSpeedTransformerLayer、按四类参数配置内核、通过内存优化开关放大 batch size，并最终用 --deepspeed_transformer_kernel 启动训练，是一份可直接照抄复用的实战指南。

Transformer Kernel 是什么

Transformer Kernel 是 DeepSpeed 为 Transformer 层专门编写的一套高性能 CUDA 算子，它将普通 PyTorch Transformer 层中分散的 GEMM、Softmax、Dropout、LayerNorm、GELU、残差连接等子操作融合进专门的核函数中，从而在单卡上提升训练吞吐、在多卡扩展时保持良好线性度。

从当前仓库源码可以看到这套实现确实是一等公民，而非简单的封装：

• Python 层的 Layer 定义与 autograd 封装位于 deepspeed/ops/transformer/transformer.py，其中 DeepSpeedTransformerFunction 直接以 torch.autograd.Function 方式自定义了前向（forward，第 145 行起）与反向（backward，第 236 行起）；
• CUDA 侧的实现源码位于 csrc/transformer，构建清单在 op_builder/transformer.py 中可见，共包含 8 个编译单元：ds_transformer_cuda.cpp、cublas_wrappers.cu、transform_kernels.cu、gelu_kernels.cu、dropout_kernels.cu、normalize_kernels.cu、softmax_kernels.cu、general_kernels.cu；
• 反向传播时同样调用 CUDA 侧的 backward_fp16/backward_fp32，一个前向/反向周期内即完成整层全部计算，避免了多次 kernel 启动与中间张量的反复读写。

这一设计与原文档给出的动机完全一致：让 Transformer 训练在性能上足够高效，使研究者能在合理时间内探索不同模型规模与超参数组合。

使用前提与注意点

环境前提

使用 Transformer Kernel 前，请先按照 Getting Started 快速开始教程 把 DeepSpeed 完整集成进你的训练脚本（即能够通过 deepspeed 启动器运行带 ZeRO/优化器配置的训练）。

兼容性提醒

原文档给出了明确警告，请务必遵守：

目前 DeepSpeed Transformer Kernel 不支持 Sparse Attention。若需使用 Sparse Attention，必须关闭 Transformer Kernel！

第一步：把 Kernel 集成进顶层模型

将 Transformer Kernel 集成进模型的过程，就是把原来 PyTorch 的 TransformerEncoderLayer 替换为 DeepSpeedTransformerLayer。

下面以 Pre-LN BERT-Large 配置为例：共 24 层、hidden size 1024、序列长度 128、batch size 64。同一份 layer 规格通过 copy.deepcopy 复制 num_hidden_layers 份，放入 nn.ModuleList，每份会拿到不同的内部 layer_id。

先实例化配置，再创建 24 个 DeepSpeedTransformerLayer：

from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig

config = DeepSpeedTransformerConfig(batch_size=64,
                                    max_seq_length=128,
                                    hidden_size=1024,
                                    heads=16,
                                    attn_dropout_ratio=0.1,
                                    hidden_dropout_ratio=0.1,
                                    num_hidden_layers=24,
                                    initializer_range=0.02,
                                    local_rank=0,
                                    seed=1234,
                                    fp16=True,
                                    pre_layer_norm=True,
                                    attn_dropout_checkpoint=False,
                                    normalize_invertible=False,
                                    gelu_checkpoint=False)
self.layer = nn.ModuleList([
    copy.deepcopy(DeepSpeedTransformerLayer(config))
    for _ in range(config.num_hidden_layers)
])

代码要点说明：

• 导入路径来自仓库导出接口，tests/unit/ops/accelerators/test_accelerator_forward.py 与 test_accelerator_backward.py 中也是 from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig；
• 每个 DeepSpeedTransformerLayer 在被实例化时，会通过类级静态变量 layer_id 自动分配从 0 递增的编号（见 transformer.py），CUDA 侧按 layer_id 对应维护各自的 buffer；
• local_rank >= 0 时构造器会主动 get_accelerator().set_device(local_rank)（见 transformer.py），保证 kernel 在正确的设备上初始化。

第二步：理解全部配置参数

配置类在源码中为 DeepSpeedTransformerConfig，继承自 TransformerConfig 并做了大量扩展。原文档将参数划分为四类，下文逐类给出说明，并补充源码中的默认值与额外参数。

1. 通用配置参数（General Configuration）

参数	含义
batch_size	每张 GPU 上运行 kernel 的 micro-batch 大小
max_seq_length	训练所用模型的最大序列长度
hidden_size	Transformer 层的隐藏维度
heads	self-attention 的头数
attn_dropout_ratio	attention 输出上的 dropout 比例
hidden_dropout_ratio	Transformer 输出（FFN 之后）上的 dropout 比例
num_hidden_layers	Transformer 层总数
pre_layer_norm	选择 Pre-LN 还是 Post-LN 架构

2. 环境参数（Environment Parameters）

参数	含义
local_rank	运行该 kernel 的当前 GPU 编号
seed	dropout 层的随机种子
fp16	是否启用半精度计算
initializer_range	BERT 初始化范围

3. 高性能优化开关

参数	含义与注意点
stochastic_mode	开启后平均可提升约 2% 的训练速度（数据来源：原教程说明）。该模式带有一定程度的非确定性，多次运行结果会有差异。经验上，BERT 这类预训练任务开启后不受影响、仍能达到高精度；但微调等下游任务建议关闭，以保证结果可复现

从源码看，stochastic_mode=True 时会加载并使用一套独立编译的随机化内核：构建类 op_builder/stochastic_transformer.py 定义了 DS_BUILD_STOCHASTIC_TRANSFORMER 环境变量，并在 nvcc 编译参数中加入 -D__STOCHASTIC_MODE__；运行时分发逻辑见 transformer.py 中 stochastic_transformer_cuda_module if config.stochastic_mode else transformer_cuda_module 的切换。

4. 内存优化开关

参数	含义
attn_dropout_checkpoint	对 attention dropout 结果做 checkpointing 以省内存
normalize_invertible	启用可逆 LayerNorm 执行（丢弃其输入激活）
gelu_checkpoint	对 GELU 激活输出做 checkpointing 以省内存

源码中出现的其它参数（补充）

阅读 DeepSpeedTransformerConfig 的构造器，可以发现一些默认值，便于你按需调整：

• 默认值：fp16=False、pre_layer_norm=True、normalize_invertible=False、gelu_checkpoint=False、attn_dropout_checkpoint=False、stochastic_mode=False；
• intermediate_size：FFN 中间维度，缺省（<= 0）时自动取 4 * hidden_size；
• layer_norm_eps=1e-12：LayerNorm 的 epsilon；
• adjust_init_range=True：为 True 时会对 self-attention 输出与 FFN 输出的权重初始化做残差路径累积修正，即 output_std = initializer_range / sqrt(2.0 * num_layers)（实现见 init_transformer_weights）；
• return_tuple=False：控制前向结果是否以元组接口返回；
• training=True：区分训练/推理模式。

第三步：内存优化开关的作用与搭配策略

Transformer Kernel 内置了多种省内存技巧，它们在层内的不同位置发挥作用，并全部以配置开关形式暴露。开启这些开关后通常能支撑更大的 batch size；尽管个别技术会以少量计算换内存，但整体收益是正的——更大的 batch size 提高了端到端训练效率。原文档对三种机制的原理说明如下：

• normalize_invertible（可逆 LayerNorm）：强制 kernel 丢弃传给 Transformer normalize 层的输入激活。之所以可以这样做，是因为 kernel 内部实现了一种优化——只需利用输出激活即可同时算出参数梯度与本层的输入梯度，无需保留输入激活。
• attn_dropout_checkpoint 与 gelu_checkpoint（checkpointing）：丢弃 Transformer 层内 attention dropout 与 GELU 两处的输入，从而省下一大块激活显存。根据项目侧的性能分析，这两个子模块重新计算（rematerialize）的性能开销可忽略不计，而由此换来更大 batch size 带来的收益足以覆盖这部分成本。

反向传播的实现（见 transformer.py 的 backward）验证了上述机制：开启 attn_dropout_checkpoint 后，ctx_bufB（dropout 前的 buffer）不再被保存，backward 中用 soft_inp 重新计算；开启 gelu_checkpoint 后不再保存 gelu_inp，backward 用 ff2_inp 重算；开启 normalize_invertible 后不再保存 normalize 层的输入 input/add_res，梯度完全由输出激活推导。

BERT-Large × V100 32GB 开关速查表

下表出自原文档，给出了在 NVIDIA V100 32GB 显存上跑 BERT-Large、面对不同 micro-batch 与序列长度时应开启的内存优化开关：

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

使用方法是按表格“从下往上/从右往左”叠加开关：序列长度 512 时，先用 attn_dropout_checkpoint 支撑到 micro-batch 16；若 batch 还要更大，再依次开启 normalize_invertible 与 gelu_checkpoint。序列长度 128 时显存压力小得多，micro-batch 超过 80 后按 normalize_invertible → attn_dropout_checkpoint → gelu_checkpoint 的顺序逐步开启即可。该表格基于 V100-32GB 实验环境，若显存规格不同，可把它当作开关组合顺序的参考基准。

第四步：用启动器启用 Transformer Kernel

启用自定义内核时，唯一要做的改动是在启动命令中传入 --deepspeed_transformer_kernel。下面是在 BERT 预训练任务中与其它参数一起传给 deepspeed 启动器的示例命令（原文档命令照录）：

deepspeed deepspeed_train.py \
--cf bert_large_lamb.json \
--max_seq_length 512 \
--print_steps 100 \
--deepspeed \
--deepspeed_transformer_kernel \
--deepspeed_config deepspeed_bsz32K_lamb_config_seq512.json \
--rewarmup \
--lr_schedule "EE" \
--lr_offset 0.0 \
--attention_dropout_checkpoint \
--load_training_checkpoint ${CHECKPOINT_BASE_PATH} \
--load_checkpoint_id ${CHECKPOINT_EPOCH150_NAME}

参数说明：

• --deepspeed_transformer_kernel：核心开关，命令解析后在代码中把模型的 Transformer 层切换为 DeepSpeed Transformer Kernel；
• --deepspeed_config：DeepSpeed ZeRO 等特性的 JSON 配置；
• --attention_dropout_checkpoint：对应配置类里的 attn_dropout_checkpoint。上例中用它来支撑 seq 512 下每卡 micro-batch 16；需要更大 batch 时可继续叠加其余内存优化开关；
• 其余 --lr_*、--load_* 等是示例 BERT 训练脚本自身的学习率与断点续训参数，实际使用时替换为你自己训练脚本的命令行。

关于 --deepspeed_transformer_kernel 是如何落地的，可参考同仓库的 BERT 预训练教程：训练脚本用 parser.add_argument('--deepspeed_transformer_kernel', ...) 注册该布尔参数，随后在构建模型时将 CLI 上的 attention_dropout_checkpoint 等选项透传进 DeepSpeedTransformerConfig（如 attn_dropout_checkpoint=args.attention_dropout_checkpoint），最终实例化 DeepSpeedTransformerLayer 替换原生层。

构建相关环境变量

Transformer Kernel 属于需要编译的 CUDA 算子，可用以下环境变量控制其构建（默认在首次使用时会 JIT 构建）：

• DS_BUILD_TRANSFORMER=1：编译标准 Transformer Kernel（见 op_builder/transformer.py）；
• DS_BUILD_STOCHASTIC_TRANSFORMER=1：额外编译随机化（stochastic）版本（见 op_builder/stochastic_transformer.py）。

底层调用链与验证方式

运行时行为

训练模式下，DeepSpeedTransformerLayer.forward（transformer.py）会先把当前是否处于梯度使能状态写入 config，再调用 DeepSpeedTransformerFunction.apply，把全部权重参数（QKV、Attention 输出、两层 FFN、两层 LayerNorm 的权重与偏置）一次性传给自定义 CUDA Function；kernel 在 fp16/fp32 两种精度路径间按 config.fp16 分发，标准/随机化两种模块间按 config.stochastic_mode 分发。

单元测试佐证

仓库在 tests/unit/ops/accelerators/test_accelerator_forward.py 与 tests/unit/ops/accelerators/test_accelerator_backward.py 中提供了可直接参考的用法：

• 与教程代码一致地通过 copy.deepcopy(DeepSpeedTransformerLayer(...)) 构造多层；
• 前向测试中会构造 DeepSpeedTransformerConfig 并设置 ds_config.stochastic_mode = True 等开关，覆盖不同配置组合下的 kernel 前向/反向正确性。

如果你要验证自己的集成是否正确，可以对照这些测试中的层构造方式与 mask/输入形状约定来排查。

总结与最佳实践清单

• 主开关：模型接入用 DeepSpeedTransformerLayer，启动训练加 --deepspeed_transformer_kernel；
• 显存吃紧时：按 attn_dropout_checkpoint → normalize_invertible → gelu_checkpoint 的顺序开启，参考上文的 V100-32GB 速查表；
• 精度与复现：预训练可开 stochastic_mode 提速约 2%；微调等对复现有要求的任务请关闭；
• 架构一致性：pre_layer_norm 必须与模型实际的 LN 位置（Pre-LN/Post-LN）一致，fp16 需与整体混合精度策略匹配；
• 组合限制：Sparse Attention 与 Transformer Kernel 不可同时使用；
• 更完整的端到端 BERT 预训练配置、序列长度/批大小适配案例与性能评估，可继续阅读 BERT 预训练教程，历史性能数据与发布说明见仓库博客 fastest-bert-training。