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/DeepSpeedTransformer 层几乎出现在所有现代序列建模模型中其训练效率直接决定了 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 个DeepSpeedTransformerLayerfrom deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig config DeepSpeedTransformerConfig(batch_size64, max_seq_length128, hidden_size1024, heads16, attn_dropout_ratio0.1, hidden_dropout_ratio0.1, num_hidden_layers24, initializer_range0.02, local_rank0, seed1234, fp16True, pre_layer_normTrue, attn_dropout_checkpointFalse, normalize_invertibleFalse, gelu_checkpointFalse) 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.pyCUDA 侧按layer_id对应维护各自的 bufferlocal_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_sizeTransformer 层的隐藏维度headsself-attention 的头数attn_dropout_ratioattention 输出上的 dropout 比例hidden_dropout_ratioTransformer 输出FFN 之后上的 dropout 比例num_hidden_layersTransformer 层总数pre_layer_norm选择 Pre-LN 还是 Post-LN 架构2. 环境参数Environment Parameters参数含义local_rank运行该 kernel 的当前 GPU 编号seeddropout 层的随机种子fp16是否启用半精度计算initializer_rangeBERT 初始化范围3. 高性能优化开关参数含义与注意点stochastic_mode开启后平均可提升约 2% 的训练速度数据来源原教程说明。该模式带有一定程度的非确定性多次运行结果会有差异。经验上BERT 这类预训练任务开启后不受影响、仍能达到高精度但微调等下游任务建议关闭以保证结果可复现从源码看stochastic_modeTrue时会加载并使用一套独立编译的随机化内核构建类 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 的构造器可以发现一些默认值便于你按需调整默认值fp16False、pre_layer_normTrue、normalize_invertibleFalse、gelu_checkpointFalse、attn_dropout_checkpointFalse、stochastic_modeFalseintermediate_sizeFFN 中间维度缺省 0时自动取4 * hidden_sizelayer_norm_eps1e-12LayerNorm 的 epsilonadjust_init_rangeTrue为 True 时会对 self-attention 输出与 FFN 输出的权重初始化做残差路径累积修正即output_std initializer_range / sqrt(2.0 * num_layers)实现见 init_transformer_weightsreturn_tupleFalse控制前向结果是否以元组接口返回trainingTrue区分训练/推理模式。第三步内存优化开关的作用与搭配策略Transformer Kernel 内置了多种省内存技巧它们在层内的不同位置发挥作用并全部以配置开关形式暴露。开启这些开关后通常能支撑更大的 batch size尽管个别技术会以少量计算换内存但整体收益是正的——更大的 batch size 提高了端到端训练效率。原文档对三种机制的原理说明如下normalize_invertible可逆 LayerNorm强制 kernel 丢弃传给 Transformer normalize 层的输入激活。之所以可以这样做是因为 kernel 内部实现了一种优化——只需利用输出激活即可同时算出参数梯度与本层的输入梯度无需保留输入激活。attn_dropout_checkpoint与gelu_checkpointcheckpointing丢弃 Transformer 层内 attention dropout 与 GELU 两处的输入从而省下一大块激活显存。根据项目侧的性能分析这两个子模块重新计算rematerialize的性能开销可忽略不计而由此换来更大 batch size 带来的收益足以覆盖这部分成本。反向传播的实现见 transformer.py 的 backward验证了上述机制开启attn_dropout_checkpoint后ctx_bufBdropout 前的 buffer不再被保存backward 中用soft_inp重新计算开启gelu_checkpoint后不再保存gelu_inpbackward 用ff2_inp重算开启normalize_invertible后不再保存 normalize 层的输入input/add_res梯度完全由输出激活推导。BERT-Large × V100 32GB 开关速查表下表出自原文档给出了在 NVIDIA V100 32GB 显存上跑 BERT-Large、面对不同 micro-batch 与序列长度时应开启的内存优化开关Micro-batch size128 sequence-length512 sequence-length 12-attn_dropout_checkpoint 16-normalize_invertible,gelu_checkpoint 80normalize_invertibleOOM 112attn_dropout_checkpointOOM 128gelu_checkpointOOM使用方法是按表格“从下往上/从右往左”叠加开关序列长度 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_configDeepSpeed 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_checkpointargs.attention_dropout_checkpoint最终实例化DeepSpeedTransformerLayer替换原生层。构建相关环境变量Transformer Kernel 属于需要编译的 CUDA 算子可用以下环境变量控制其构建默认在首次使用时会 JIT 构建DS_BUILD_TRANSFORMER1编译标准 Transformer Kernel见 op_builder/transformer.pyDS_BUILD_STOCHASTIC_TRANSFORMER1额外编译随机化stochastic版本见 op_builder/stochastic_transformer.py。底层调用链与验证方式运行时行为训练模式下DeepSpeedTransformerLayer.forwardtransformer.py会先把当前是否处于梯度使能状态写入config再调用DeepSpeedTransformerFunction.apply把全部权重参数QKV、Attention 输出、两层 FFN、两层 LayerNorm 的权重与偏置一次性传给自定义 CUDA Functionkernel 在 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/10 4:50:39 · 阅读 1.2k 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/DeepSpeedTransformer 层几乎出现在所有现代序列建模模型中其训练效率直接决定了 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 个DeepSpeedTransformerLayerfrom deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig config DeepSpeedTransformerConfig(batch_size64, max_seq_length128, hidden_size1024, heads16, attn_dropout_ratio0.1, hidden_dropout_ratio0.1, num_hidden_layers24, initializer_range0.02, local_rank0, seed1234, fp16True, pre_layer_normTrue, attn_dropout_checkpointFalse, normalize_invertibleFalse, gelu_checkpointFalse) 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.pyCUDA 侧按layer_id对应维护各自的 bufferlocal_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_sizeTransformer 层的隐藏维度headsself-attention 的头数attn_dropout_ratioattention 输出上的 dropout 比例hidden_dropout_ratioTransformer 输出FFN 之后上的 dropout 比例num_hidden_layersTransformer 层总数pre_layer_norm选择 Pre-LN 还是 Post-LN 架构2. 环境参数Environment Parameters参数含义local_rank运行该 kernel 的当前 GPU 编号seeddropout 层的随机种子fp16是否启用半精度计算initializer_rangeBERT 初始化范围3. 高性能优化开关参数含义与注意点stochastic_mode开启后平均可提升约 2% 的训练速度数据来源原教程说明。该模式带有一定程度的非确定性多次运行结果会有差异。经验上BERT 这类预训练任务开启后不受影响、仍能达到高精度但微调等下游任务建议关闭以保证结果可复现从源码看stochastic_modeTrue时会加载并使用一套独立编译的随机化内核构建类 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 的构造器可以发现一些默认值便于你按需调整默认值fp16False、pre_layer_normTrue、normalize_invertibleFalse、gelu_checkpointFalse、attn_dropout_checkpointFalse、stochastic_modeFalseintermediate_sizeFFN 中间维度缺省 0时自动取4 * hidden_sizelayer_norm_eps1e-12LayerNorm 的 epsilonadjust_init_rangeTrue为 True 时会对 self-attention 输出与 FFN 输出的权重初始化做残差路径累积修正即output_std initializer_range / sqrt(2.0 * num_layers)实现见 init_transformer_weightsreturn_tupleFalse控制前向结果是否以元组接口返回trainingTrue区分训练/推理模式。第三步内存优化开关的作用与搭配策略Transformer Kernel 内置了多种省内存技巧它们在层内的不同位置发挥作用并全部以配置开关形式暴露。开启这些开关后通常能支撑更大的 batch size尽管个别技术会以少量计算换内存但整体收益是正的——更大的 batch size 提高了端到端训练效率。原文档对三种机制的原理说明如下normalize_invertible可逆 LayerNorm强制 kernel 丢弃传给 Transformer normalize 层的输入激活。之所以可以这样做是因为 kernel 内部实现了一种优化——只需利用输出激活即可同时算出参数梯度与本层的输入梯度无需保留输入激活。attn_dropout_checkpoint与gelu_checkpointcheckpointing丢弃 Transformer 层内 attention dropout 与 GELU 两处的输入从而省下一大块激活显存。根据项目侧的性能分析这两个子模块重新计算rematerialize的性能开销可忽略不计而由此换来更大 batch size 带来的收益足以覆盖这部分成本。反向传播的实现见 transformer.py 的 backward验证了上述机制开启attn_dropout_checkpoint后ctx_bufBdropout 前的 buffer不再被保存backward 中用soft_inp重新计算开启gelu_checkpoint后不再保存gelu_inpbackward 用ff2_inp重算开启normalize_invertible后不再保存 normalize 层的输入input/add_res梯度完全由输出激活推导。BERT-Large × V100 32GB 开关速查表下表出自原文档给出了在 NVIDIA V100 32GB 显存上跑 BERT-Large、面对不同 micro-batch 与序列长度时应开启的内存优化开关Micro-batch size128 sequence-length512 sequence-length 12-attn_dropout_checkpoint 16-normalize_invertible,gelu_checkpoint 80normalize_invertibleOOM 112attn_dropout_checkpointOOM 128gelu_checkpointOOM使用方法是按表格“从下往上/从右往左”叠加开关序列长度 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_configDeepSpeed 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_checkpointargs.attention_dropout_checkpoint最终实例化DeepSpeedTransformerLayer替换原生层。构建相关环境变量Transformer Kernel 属于需要编译的 CUDA 算子可用以下环境变量控制其构建默认在首次使用时会 JIT 构建DS_BUILD_TRANSFORMER1编译标准 Transformer Kernel见 op_builder/transformer.pyDS_BUILD_STOCHASTIC_TRANSFORMER1额外编译随机化stochastic版本见 op_builder/stochastic_transformer.py。底层调用链与验证方式运行时行为训练模式下DeepSpeedTransformerLayer.forwardtransformer.py会先把当前是否处于梯度使能状态写入config再调用DeepSpeedTransformerFunction.apply把全部权重参数QKV、Attention 输出、两层 FFN、两层 LayerNorm 的权重与偏置一次性传给自定义 CUDA Functionkernel 在 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考