DeepSpeed Progressive Layer DroppingPLD实战以渐进式层丢弃加速 Transformer 预训练【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeedProgressive Layer DroppingPLD渐进式层丢弃是 DeepSpeed 提出的压缩训练技术它通过在 Transformer 预训练中按渐进式调度稀疏地更新 Transformer Block在几乎不牺牲下游任务精度的前提下显著降低训练开销。本文以 DeepSpeed 仓库中官方发布与技术教程为主线结合 runtime 源码 与 单元测试系统讲解 PLD 的原理、θ/γ 超参的调度公式、DeepSpeed 配置启用方法并给出完整的 BERT 预训练与 GLUE 微调实操流程与预期结果帮助你在自己的 Transformer 预训练任务中复现同样的加速效果。PLD 的动机Transformer 预训练的高昂代价Transformer 类网络如 BERT的预训练往往伴随着整体计算开销过大的问题。PLD 正是针对这一痛点提出研究者分析 Transformer 训练的动态过程与稳定性后提出在训练中对 Transformer Block 进行稀疏更新——即每个 mini-batch 只更新部分层、跳过其余层的前向/反向计算。其丢弃率并非固定不变而是沿着时间training step与模型深度两个维度平滑增长的渐进式丢弃调度。根据官方发布见 news 原文与技术报告仓库文档中引用的结论PLD 可在下游任务取得相近精度的前提下让预训练提速约 2.5 倍在训练相同数量样本时训练速度提升约 24%该加速不依赖额外的硬件资源即不增加显存或卡数开销。需要说明的是这两个数字来自上述官方技术报告在特定 BERT 预训练设定下的实验结论实际加速比会随模型规模、批大小与硬件配置而变化可作为复现基线而非普适承诺。核心机制theta、gamma 与丢弃率调度公式PLD 的运行时实现非常精简全部逻辑集中在 deepspeed/runtime/progressive_layer_drop.py 的ProgressiveLayerDrop类中class ProgressiveLayerDrop(object): def __init__(self, theta0.5, gamma0.001): self.theta theta self.gamma gamma self.current_theta 1.0 def get_state(self): kwargs {progressive_layer_drop: True, pld_theta: self.get_theta()} return kwargs def get_theta(self): return self.current_theta def update_state(self, global_step): def _prob(x, gamma, p): return (1. - p) * np.exp(-gamma * x) p self.current_theta _prob(global_step, self.gamma, self.theta)两个核心超参的含义与类注释一致参数含义默认值源码官方推荐教程/发布theta权衡训练速度与模型鲁棒性的目标丢弃率值越低训练越快1.0见 constants.py 中PLD_THETA_DEFAULT0.5gamma控制丢弃率从 1.0 指数衰减到theta的速度0.0010.001调度公式为theta(t) (1 - theta) * exp(-gamma * t) theta训练初期current_theta从1.0起步即暂时不丢弃层保证训练稳定随全局步数t的增长按指数曲线衰减最终平滑逼近目标值theta。这就是文档中所说的沿时间维度渐进的由来——模型先以完整结构稳定预热的早期阶段再逐步提高稀疏度。丢弃信号如何注入模型前向get_state()返回的{progressive_layer_drop: True, pld_theta: ...}会被 DeepSpeed 引擎注入模型前向调用见 engine.py 的 forward 调用处训练模式下若self.progressive_layer_drop存在则用kwargs.update(...)把progressive_layer_dropTrue与当前pld_theta传给模型。因此启用 PLD 要求模型的前向函数接收**kwargs并依据pld_theta决定每层是否被跳过沿深度维度的随机丢弃。仓库测试中的PLD_SimpleModel即按此契约实现def forward(self, x, y, **kwargs): pld kwargs.get(progressive_layer_drop, False) theta kwargs.get(pld_theta, 1.0) ...见 tests/unit/simple_model.py。在 DeepSpeed 中启用 PLD启用 PLD 需要同时在客户端脚本与DeepSpeed 引擎两侧打开开关。第一步在客户端脚本模型侧添加命令行参数--progressive_layer_drop第二步在 DeepSpeed JSON 配置中加入progressive_layer_drop配置字典{ ... progressive_layer_drop: { enabled: true, theta: 0.5, gamma: 0.001 } }官方实验表明theta0.5、gamma0.001的组合效果良好可作为推荐起点。配置如何被解析从源码看配置解析位于 deepspeed/runtime/config.pyget_pld_enabled读取字典中的enabled键默认Falseget_pld_params取出并剥离enabled后的剩余参数作为 PLD 参数集随后DeepSpeedConfig在 config.py#L883-L884 保存pld_enabled与pld_params。对应键名常量定义在 constants.pyenabled默认False、theta默认1.0、gamma默认0.001。引擎在 engine.py 检测到pld_enabled()为真时通过_configure_progressive_layer_drop()engine.py#L2658-L2661实例化调度器并在每个全局步结束时调用update_state(self.global_steps)engine.py#L3556-L3557推进current_theta。训练过程中可通过engine.get_pld_theta()engine.py#L3742-L3746随时查询当前丢弃率。第三步确认运行时日志配置成功后引擎会打印如下日志仅 Rank 0[INFO] [logging.py:60:log_dist] [Rank 0] Enabled progressive layer dropping (theta 0.5)该日志由ProgressiveLayerDrop.__init__中的log_dist(..., ranks[0])产生可作为启用成功的直接证据。实战一用 DeepSpeed PLD 预训练 BERT数据准备预训练数据Wikipedia 与 BookCorpus与 BERT 原文类似的下载与预处理流程请参照仓库内 BERT 预训练教程其中包含详细的数据下载与处理说明。启动脚本预训练主逻辑位于已改造为使用 DeepSpeed 的deepspeed_train.py中ds_train_bert_progressive_layer_drop_bsz4k_seq128.sh则是带 PLD 的启动脚本运行方式bash ds_train_bert_progressive_layer_drop_bsz4k_seq128.sh脚本中的大多数参数与标准 BERT 预训练教程一致区别即上文所述的--progressive_layer_drop客户端参数与配置中的progressive_layer_drop字典。完整 DeepSpeed 配置下方为官方用于 BERTPLD 预训练的完整配置同时涵盖批大小、优化器、学习率、序列长度等{ train_batch_size: 4096, train_micro_batch_size_per_gpu: 16, steps_per_print: 1000, prescale_gradients: true, gradient_predivide_factor: 8, optimizer: { type: Adam, params: { lr: 1e-3, weight_decay: 0.01, bias_correction: false } }, gradient_clipping: 1.0, wall_clock_breakdown: false, fp16: { enabled: true, loss_scale: 0 }, progressive_layer_drop: { enabled: true, theta: 0.5, gamma: 0.001 } }该配置面向64 × 32GB V100 GPU每卡 micro batch 为 16通过梯度累积把有效批大小凑到 4096。若 GPU 显存较小需调小train_micro_batch_size_per_gpu若 GPU 数量更多则可增大train_batch_size进一步提升吞吐。预训练超参数总览参数取值Effective batch size4KTrain micro batch size per GPU16OptimizerAdamPeak learning rate1e-3Sequence-length128Learning rate schedulerWarmup linear decay expWarmup ratio0.02Decay rate0.99Decay step1000Weight decay0.01Gradient clipping1.0关于 PreLayerNorm 的说明官方提示DeepSpeed 已将PreLayerNorm作为训练 BERT 的默认方式可避免梯度消失、稳定优化并带来性能收益详见仓库fastest BERT training发布文章。因此可切换的 Transformer Block 直接构建在 PreLayerNorm BERT 之上——这正是 PLD 能沿深度维度安全跳过层的前提之一先做层归一化可让网络在部分层被随机跳过的情形下仍保持前向/反向传播的数值稳定。实战二用 PLD 预训练模型微调 GLUEGLUEGeneral Language Understanding Evaluation benchmark是句子级与句对级自然语言理解任务的集合涵盖问答、情感分析、文本蕴含等设计上鼓励跨任务的知识迁移与样本高效学习。准备 GLUE 数据可使用社区提供的辅助脚本下载全部 GLUE 数据教程原文给出的 gist 脚本数据下载后将其放置到默认路径/data/GlueData即可。指定 BERT 模型配置微调主逻辑位于已改造为使用 DeepSpeed 的run_glue_classifier_bert_base.py。微调前需把其中的bert_model_config修改为与 PLD 预训练模型一致的配置bert_model_config { vocab_size_or_config_json_file: 119547, hidden_size: 768, num_hidden_layers: 12, num_attention_heads: 12, intermediate_size: 3072, hidden_act: gelu, hidden_dropout_prob: 0.1, attention_probs_dropout_prob: 0.1, max_position_embeddings: 512, type_vocab_size: 2, initializer_range: 0.02 }即标准的 BERT-Base 结构12 层、hidden 768、12 头、FFN 中间 3072。加载 DeepSpeed checkpoint脚本中已内置 DeepSpeed 风格 checkpoint 的加载方式核心一行model.load_state_dict(checkpoint_state_dict[module], strictFalse)strictFalse允许状态字典键不完全一一对应例如只加载预训练部分权重这在复用预训练模型做分类头微调时很关键。启动微调run_glue_classifier_bert_base.sh脚本封装了微调所需的超参数调用方式bash run_glue_bert_base_finetune.sh [task] [batch size] [learning rate] [number of epochs] [job name] [checkpoint path]一个示例在 MNLI 任务上微调学习率 3e-5bash run_glue_bert_base_finetune.sh MNLI 32 3e-5 5 fine_tune_MNLI deepspeed_checkpoint.pt预期微调结果微调日志输出在logs目录。下表对比了原始 BERT-Base、DeepSpeed 复现的 BERT-Base 与PLDPLD 预训练后微调在各 GLUE 任务上的结果Lr行是取得对应结果所用学习率RTEMRPCSTS-BCoLASST-2QNLIQQPMNLI-m/mmGLUEMetricsAcc.F1/Acc.PCC/SCCAcc.Acc.Acc.F1/Acc.Acc.Bert_{base} (original)66.488.9/84.887.1/89.252.193.590.571.2/89.284.6/83.480.7Bert_{base} (Our impl)67.888.0/86.089.5/89.252.591.287.189.0/90.682.5/83.482.1PLD69.386.6/84.390.0/89.655.891.690.789.6/91.284.1/83.882.9Lr7e-59e-57e-55e-57e-59e-52e-43e-5可以看到PLD 预训练模型在 GLUE 平均分82.9上不仅没有因跳过层而受损反而略高于两个对照基线这正是压缩训练不牺牲精度的核心论据。仓库内的验证测试PLD 的调度逻辑有专门的单元测试覆盖见 tests/unit/runtime/test_pld.pytest_pld_schedule针对theta ∈ {0, 0.1, 0.9, 1.0}逐次调用update_state(i)后断言current_theta精确等于(1 - theta) * exp(-gamma * i) theta直接锁定调度公式的正确性TestPLDModel::test_pld_model在单卡分布式环境下用完整 DeepSpeed 配置含progressive_layer_drop字典初始化PLD_SimpleModel并真实训练 50 个样本每一步校验engine.get_pld_theta()与公式一致验证了引擎在每个 global step 后正确推进调度状态TestNonPLDModel::test_non_pld_model使用不接受**kwargs的普通SimpleModel前向会抛出TypeError——这从反面印证了模型必须实现 PLD 感知的前向签名这一接入前提。小结从本文可以看出PLD 是一套思路简单、工程闭环的训练加速方案思想上用一条指数衰减曲线在时间维度上渐进引入层丢弃实现在深度维度上的稀疏更新工程上则仅需一个 JSON 配置块加一个客户端参数DeepSpeed 引擎负责调度推进与丢弃信号注入模型侧只需让前向接收pld_theta。再配合仓库内完整的 BERT 预训练教程、GLUE 微调基线表与单元测试你完全可以按官方设定在自有 Transformer 预训练任务上快速复现同精度更高吞吐的效果。进一步资料可继续阅读仓库内完整的 Progressive Layer Dropping 教程。【免费下载链接】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），仅供参考