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

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

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

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）

2. 环境参数（Environment Parameters）

3. 高性能优化开关

从源码看，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. 内存优化开关

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

阅读 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 与序列长度时应开启的内存优化开关：

使用方法是按表格“从下往上/从右往左”叠加开关：序列长度 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。

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed