DeepSpeed Progressive Layer Dropping（PLD）实战：以渐进式层丢弃加速 Transformer 预训练

Progressive Layer Dropping（PLD，渐进式层丢弃）是 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, theta=0.5, gamma=0.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_DEFAULT）	0.5
gamma	控制丢弃率从 1.0 指数衰减到 theta 的速度	0.001	0.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_drop=True 与当前 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
  }
}

官方实验表明 theta=0.5、gamma=0.001 的组合效果良好，可作为推荐起点。

配置如何被解析

从源码看，配置解析位于 deepspeed/runtime/config.py：get_pld_enabled 读取字典中的 enabled 键（默认 False），get_pld_params 取出并剥离 enabled 后的剩余参数作为 PLD 参数集；随后 DeepSpeedConfig 在 config.py#L883-L884 保存 pld_enabled 与 pld_params。对应键名常量定义在 constants.py：enabled（默认 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 配置

下方为官方用于 BERT+PLD 预训练的完整配置，同时涵盖批大小、优化器、学习率、序列长度等：

{
  "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 size	4K
Train micro batch size per GPU	16
Optimizer	Adam
Peak learning rate	1e-3
Sequence-length	128
Learning rate scheduler	Warmup linear decay exp
Warmup ratio	0.02
Decay rate	0.99
Decay step	1000
Weight decay	0.01
Gradient clipping	1.0

关于 PreLayerNorm 的说明

官方提示：DeepSpeed 已将 PreLayerNorm 作为训练 BERT 的默认方式（可避免梯度消失、稳定优化并带来性能收益，详见仓库"fastest BERT training"发布文章）。因此可切换的 Transformer Block 直接构建在 PreLayerNorm BERT 之上——这正是 PLD 能沿深度维度安全跳过层的前提之一：先做层归一化可让网络在部分层被随机跳过的情形下仍保持前向/反向传播的数值稳定。

实战二：用 PLD 预训练模型微调 GLUE

GLUE（General 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'], strict=False)

strict=False 允许状态字典键不完全一一对应（例如只加载预训练部分权重），这在复用预训练模型做分类头微调时很关键。

启动微调

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-5）：

bash run_glue_bert_base_finetune.sh MNLI 32 3e-5 5 "fine_tune_MNLI" deepspeed_checkpoint.pt

预期微调结果

微调日志输出在 logs 目录。下表对比了原始 BERT-Base、DeepSpeed 复现的 BERT-Base 与 PLD（PLD 预训练后微调）在各 GLUE 任务上的结果，Lr 行是取得对应结果所用学习率：

	RTE	MRPC	STS-B	CoLA	SST-2	QNLI	QQP	MNLI-m/mm	GLUE
Metrics	Acc.	F1/Acc.	PCC/SCC	Acc.	Acc.	Acc.	F1/Acc.	Acc.
Bert_{base} (original)	66.4	88.9/84.8	87.1/89.2	52.1	93.5	90.5	71.2/89.2	84.6/83.4	80.7
Bert_{base} (Our impl)	67.8	88.0/86.0	89.5/89.2	52.5	91.2	87.1	89.0/90.6	82.5/83.4	82.1
PLD	69.3	86.6/84.3	90.0/89.6	55.8	91.6	90.7	89.6/91.2	84.1/83.8	82.9
Lr	7e-5	9e-5	7e-5	5e-5	7e-5	9e-5	2e-4	3e-5

可以看到，PLD 预训练模型在 GLUE 平均分（82.9）上不仅没有因"跳过层"而受损，反而略高于两个对照基线，这正是"压缩训练不牺牲精度"的核心论据。

仓库内的验证测试

PLD 的调度逻辑有专门的单元测试覆盖，见 tests/unit/runtime/test_pld.py：

• test_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 教程。