PaddleNLP AdamW 优化器实战:AdamWDL 分层学习率衰减的原理、源码解析与使用方法

发布时间:2026/9/23 23:18:17
PaddleNLP AdamW 优化器实战:AdamWDL 分层学习率衰减的原理、源码解析与使用方法 人工智能大模型预训练微调LoRARLHF强化学习分布式训练【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址https://gitcode.com/gh_mirrors/pa/PaddleNLP点击查看免费下载AdamW 优化器是 Transformer 类大模型训练的标配优化器而分层学习率衰减Layer-wise Learning Rate Decay则能进一步提升深层模型与预训练模型的收敛质量。本文以 docs/zh/source/paddlenlp.ops.optimizer.AdamwOptimizer.rst 对应的 PaddleNLP 优化器 API 为核心深入讲解 PaddleNLP 在paddlenlp.ops.optimizer中实现的AdamWDLAdamW with Dynamic layer-wise lr与layerwise_lr_decay工具函数。读完本文你将掌握 AdamW 与分层学习率衰减的数学原理、AdamWDL的全部参数语义与默认值、其底层优化算子执行链路的源码实现并能在自己的训练脚本中直接落地使用。一、背景为什么 Transformer 训练需要 AdamW 与分层学习率衰减Transformer 模型尤其是 BERT 系、GPT 系等深层预训练模型在训练时通常面临两个经典问题权重衰减Weight Decay与 L2 正则的耦合问题标准 SGD 下的 L2 正则与权重衰减等价但在 Adam 这类基于自适应学习率的优化器中L2 正则会被一阶矩/二阶矩估计放大导致权重衰减行为偏离预期。AdamW 将权重衰减与梯度自适应更新解耦decoupled weight decay即衰减项直接作用于参数本身而非梯度从而得到更稳定、更可控的正则化效果。不同层对学习率敏感度不同浅层如 embedding通常学习速度慢、需要更保守的学习率深层则相反。分层学习率衰减Layer-wise Learning Rate Decay通过按层指数衰减学习率让浅层使用更小的学习率、深层使用更大的学习率从而提升收敛质量。PaddleNLP 正是围绕这两个需求在 paddlenlp/ops/optimizer/adamwdl.py 中实现了AdamWDL优化器它继承自 Paddle 官方paddle.optimizer.AdamW在标准 AdamW 的基础上叠加每层动态学习率能力默认使用layerwise_lr_decay作为动态学习率设置函数。说明docs/zh/source/paddlenlp.ops.optimizer.AdamwOptimizer.rst是一份由 Sphinxautomodule指令生成的 API 文档页其内容来源于paddlenlp.ops.optimizer模块对应优化器类的 docstring。当前仓库中该文档对应的实际实现即AdamWDL类位于 paddlenlp/ops/optimizer/adamwdl.py同系列 API 文档还包括 paddlenlp.ops.optimizer.adamwdl.rst、paddlenlp.ops.optimizer.adamw.rst 等。二、模块总览paddlenlp.ops.optimizer 里有什么先看 paddlenlp/ops/optimizer/init.py 的导出内容from .adamwdl import AdamWDL, layerwise_lr_decay from .ema import ExponentialMovingAverage from .lr import InverseSquareRootSchedule __all__ [layerwise_lr_decay, AdamWDL, ExponentialMovingAverage, InverseSquareRootSchedule]该目录下共四个源文件文件导出内容作用adamwdl.pyAdamWDL、layerwise_lr_decay带分层学习率动态设置的 AdamW 优化器及其学习率衰减函数ema.pyExponentialMovingAverage指数滑动平均EMAlr.pyInverseSquareRootSchedule反平方根学习率调度__init__.py上述全部统一导出入口使用方式非常简单from paddlenlp.ops.optimizer import AdamWDL, layerwise_lr_decay三、AdamWDL完整的参数语义与默认值AdamWDL的类 docstring见 paddlenlp/ops/optimizer/adamwdl.py#L46-L131完整定义了其数学形式与全部参数。其更新规则为t t 1 moment_1_out β1 * moment_1 (1 - β1) * grad moment_2_out β2 * moment_2 (1 - β2) * grad * grad learning_rate learning_rate * sqrt(1 - β2^t) / (1 - β1^t) param_out param - learning_rate * (moment_1 / (sqrt(moment_2) ε) λ * param)其中β1、β2分别为一阶、二阶矩估计的指数衰减率ε为数值稳定小量λ为权重衰减系数。注意公式末尾的λ * param项即解耦权重衰减——它不进入梯度而是直接作用在参数更新上。3.1 构造参数全表AdamWDL.__init__的签名与默认值源码 paddlenlp/ops/optimizer/adamwdl.py#L133-L150参数类型默认值说明learning_ratefloat / LRScheduler0.001学习率可为浮点数或学习率调度器beta1float / Tensor0.9一阶矩估计指数衰减率shape 为 [1] 的 float32 Tensor 亦可beta2float / Tensor0.999二阶矩估计指数衰减率epsilonfloat / Tensor1e-08数值稳定小量parameterslist / tupleNone待更新的参数动转静模式下必填静态模式下为 None 时更新全部参数weight_decayfloat / Tensor0.01权重衰减系数解耦形式apply_decay_param_funfunction / NoneNone回调函数仅对apply_decay_param_fun(Tensor.name) True的参数施加权重衰减grad_clipGradientClipBaseNone梯度裁剪策略如GradientClipByGlobalNorm、GradientClipByNorm、GradientClipByValueNone 表示不裁剪lazy_modeboolFalse惰性模式仅更新当前 mini-batch 中有梯度的元素可显著加速超大参数更新但语义与标准 Adam 略有差异multi_precisionboolFalse是否在权重更新时使用多精度如 fp16 训练时维护 fp32 master weightlayerwise_decayfloat1.0分层学习率衰减比率1.0 表示不衰减n_layersint12Transformer 编码器总层数set_param_lr_funfunction / Nonelayerwise_lr_decay在 Adam 算子执行前设置参数学习率的回调函数name_dictdictNone键为模型动态名model.named_parameters()所得、值为静态名结构名的映射namestrNone优化器名称一般无需设置其中layerwise_decay、n_layers、set_param_lr_fun、name_dict四个参数是AdamWDL区别于普通AdamW的核心用于实现逐层动态学习率。3.2 参数校验逻辑构造函数开头对layerwise_decay做了严格类型校验源码 paddlenlp/ops/optimizer/adamwdl.py#L151-L155if not isinstance(layerwise_decay, float) and not isinstance(layerwise_decay, paddle.framework.Variable): raise TypeError(coeff should be float or Tensor.) self.layerwise_decay layerwise_decay self.n_layers n_layers self.set_param_lr_fun partial(set_param_lr_fun, layerwise_decay, name_dict, n_layers)也就是说layerwise_decay必须是float或 PaddleVariableTensor否则直接抛出TypeError。同时set_param_lr_fun通过functools.partial预绑定layerwise_decay、name_dict、n_layers三个参数后续调用时只需传入param即可。随后super().__init__(...)将剩余参数透传给父类paddle.optimizer.AdamW因此AdamWDL天然继承 Paddle 官方 AdamW 的全部能力lazy_mode、multi_precision、grad_clip等。四、核心机制layerwise_lr_decay 分层学习率衰减4.1 数学原理分层学习率衰减即自顶向下指数衰减各层学习率。例如假设第 24 层使用学习率l分层衰减率为α则第m层的学习率为l * α^(24-m)该思想详见 1906.08237 相关工作。浅层靠近 embedding学习率最小深层学习率最大。4.2 源码实现layerwise_lr_decay的实现位于 paddlenlp/ops/optimizer/adamwdl.py#L23-L43def layerwise_lr_decay(decay_rate, name_dict, n_layers, param): ratio 1.0 static_name name_dict[param.name] if encoder.layers in static_name: idx static_name.find(encoder.layers.) layer int(static_name[idx:].split(.)[2]) ratio decay_rate ** (n_layers - layer) elif embedding in static_name: ratio decay_rate ** (n_layers 1) return ratio其逻辑可拆解为名称映射通过name_dict将参数的动态名param.name映射为静态结构名如encoder.layers.0.self_attn.q_proj.weight。编码器层若静态名包含encoder.layers.解析出层号layer计算ratio decay_rate ** (n_layers - layer)。即第 0 层学习率为基准学习率的α^n_layers最后一层layer n_layers - 1为α^1顶层结构如 pooler、head保持ratio 1.0。embedding 层若静态名包含embedding学习率缩放到α^(n_layers1)是所有层中最小的符合embedding 学习最保守的实践经验。其他结构ratio保持 1.0学习率不做缩放。最终返回的ratio会被乘到该参数当前的learning_rate上再执行 Adam 更新算子详见下一节。从源码结构看layerwise_lr_decay面向的是静态名中包含encoder.layers/embedding的 Transformer 编码器类模型如 BERT 系对于不含这些关键字的模型所有参数保持 1.0等同于普通 AdamW不会产生错误。若你的模型静态名规则不同可以自定义set_param_lr_fun回调见第六节示例。五、源码解析优化算子执行链路AdamWDL通过覆写三个关键方法实现解耦权重衰减 动态学习率的完整执行链路。5.1 _append_optimize_op动态学习率的注入点源码 paddlenlp/ops/optimizer/adamwdl.py#L182-L194def _append_optimize_op(self, block, param_and_grad): if self.set_param_lr_fun is None: return super(AdamWDL, self)._append_optimize_op(block, param_and_grad) self._append_decoupled_weight_decay(block, param_and_grad) prev_lr param_and_grad[0].optimize_attr[learning_rate] ratio self.set_param_lr_fun(param_and_grad[0]) param_and_grad[0].optimize_attr[learning_rate] * ratio # execute Adam op res super(AdamWDL, self)._append_optimize_op(block, param_and_grad) param_and_grad[0].optimize_attr[learning_rate] prev_lr return res执行顺序非常清晰若set_param_lr_fun为 None退化为父类标准 AdamW 逻辑先执行解耦权重衰减_append_decoupled_weight_decay保存当前学习率prev_lr调用set_param_lr_fun(param)得到该参数的缩放系数ratio将optimize_attr[learning_rate]临时放大调用父类逻辑执行 Adam 算子此时算子使用缩放后的学习率立即恢复prev_lr避免影响其他参数。这一临时修改、用完即恢复的设计保证了每个参数都能携带各自的动态学习率进入 Adam 算子同时不会污染全局学习率状态。5.2 _append_decoupled_weight_decay解耦权重衰减实现源码 paddlenlp/ops/optimizer/adamwdl.py#L196-L243核心更新式为parameter parameter - parameter * coeff * lr关键实现细节参数过滤若_apply_decay_param_fun不为 None则只有_apply_decay_param_fun(param.name) True的参数才执行衰减学习率获取学习率为 float 时直接使用否则调用_create_param_lr为参数创建对应的学习率张量该方法必须在optimizer._create_global_learning_rate()之后调用源码注释对此有明确说明衰减系数缓存decay_coeff 1.0 - learning_rate * self._coeff并以learning_rate为键缓存在self._lr_to_coeff字典中同一学习率的参数复用计算结果避免重复计算多精度支持当multi_precisionTrue且参数为 fp16 时从self._master_weights中取出 fp32 的 master weight用decay_coeff缩放并paddle.assign写回 master weight否则直接对 fp16 参数本体做缩放设备守卫通过paddle.static.device_guard(None)保证decay_coeff计算可兼容 pipeline 等场景的设备设置。5.3 _create_optimization_pass动转静下的缓存清理源码 paddlenlp/ops/optimizer/adamwdl.py#L245-L249def _create_optimization_pass(self, parameters_and_grads): optimize_ops super(AdamWDL, self)._create_optimization_pass(parameters_and_grads) # In dygraph mode, clear _lr_to_coeff after applied gradient self._lr_to_coeff dict() return optimize_ops由于动转静dygraph模式下每步都会执行apply_gradient_lr_to_coeff缓存必须每步清空否则缓存的decay_coeff会与更新的学习率不一致源码注释明确说明该动机。此外AdamWDL还覆写了_update_param_group以支持按参数组param group更新coeff以及__str__返回Weight Decay, params: ...形式的可读描述。六、完整使用示例6.1 标准用法基于 name_dict 的分层衰减以下代码来自类 docstring 的官方示例paddlenlp/ops/optimizer/adamwdl.py#L103-L130展示了对paddle.nn.Linear应用带分层衰减的AdamWDLimport paddle from paddlenlp.ops.optimizer import AdamWDL def simple_lr_setting(decay_rate, name_dict, n_layers, param): ratio 1.0 static_name name_dict[param.name] if weight in static_name: ratio decay_rate**0.5 param.optimize_attr[learning_rate] * ratio linear paddle.nn.Linear(10, 10) name_dict dict() for n, p in linear.named_parameters(): name_dict[p.name] n inp paddle.rand([10, 10], dtypefloat32) out linear(inp) loss paddle.mean(out) adamwdl AdamWDL( learning_rate1e-4, parameterslinear.parameters(), set_param_lr_funsimple_lr_setting, layerwise_decay0.8, name_dictname_dict) loss.backward() adamwdl.step() adamwdl.clear_grad()要点拆解name_dict通过model.named_parameters()构建键为参数的动态名param.name值为结构名n自定义simple_lr_setting(decay_rate, name_dict, n_layers, param)与layerwise_lr_decay签名一致内部直接修改param.optimize_attr[learning_rate]即可AdamWDL会在_append_optimize_op中读取并应用对 Transformer 编码器模型直接使用默认的layerwise_lr_decay并传入正确的n_layers与name_dict即可无需自定义回调。6.2 典型训练循环骨架将AdamWDL接入自定义训练脚本的最小骨架import paddle from paddlenlp.ops.optimizer import AdamWDL model YourTransformerModel() # 例如 BERT 系模型 n_layers model.config.num_hidden_layers name_dict {p.name: n for n, p in model.named_parameters()} optimizer AdamWDL( learning_ratepaddle.optimizer.lr.LinearDecay(1e-4, int(1e4), 0.0), parametersmodel.parameters(), weight_decay0.01, grad_clippaddle.nn.ClipGradByGlobalNorm(1.0), layerwise_decay0.8, # 分层衰减率 n_layersn_layers, # 与实际层数保持一致 name_dictname_dict, multi_precisionTrue, # 开启 fp16 多精度更新 ) for step, batch in enumerate(dataloader): loss model(batch) loss.backward() optimizer.step() optimizer.clear_grad()注意layerwise_decay1.0默认值等价于关闭分层衰减此时AdamWDL退化为普通AdamW但代码仍可正常运行。七、优化器状态命名与 checkpoint 的关联AdamWDL继承自paddle.optimizer.AdamW在保存/加载优化器状态时会使用 Adam 系优化器的标准状态命名后缀。这一点在 PaddleNLP 的 checkpoint 转换工具 paddlenlp/trainer/utils/ckpt_converter.py 中有直接体现——该文件第 39 行定义了OPTIMIZER_STATE_NAME_SUFFIX [.moment1, .moment2, .beta1_pow_acc, .beta2_pow_acc, .master_weight]即 Adam 系优化器状态由四类张量组成状态后缀含义.moment1一阶矩估计Adam 的 momentum.moment2二阶矩估计Adam 的 variance.beta1_pow_accβ1 的幂次累积用于偏差校正.beta2_pow_accβ2 的幂次累积用于偏差校正.master_weight多精度multi_precision训练时维护的 fp32 master weightckpt_converter 在加载混合并行TP/PP/Shardingcheckpoint 时会按上述规则对优化器状态做重命名与拼接参见其rename_using_optimizer_state_order、rename_auto_parallel_state_dict等方法的注释其中明确提到adamw_optimizer_param_suffix_name_mapping与adamw_optimizer_status_name_suffix_mappings两套命名规则。这意味着若你使用AdamWDL或任一paddle.optimizer.AdamW系优化器训练其优化器状态与 PaddleNLP 的 checkpoint 转换/续训链路是天然兼容的无需额外适配。八、注意事项与最佳实践n_layers必须与模型实际层数一致layerwise_lr_decay用decay_rate ** (n_layers - layer)计算每层缩放若n_layers传错衰减幅度会整体偏移。建议从模型 config 的num_hidden_layers字段动态获取。name_dict依赖静态名规则默认衰减函数要求静态名包含encoder.layers.与embedding关键字。若模型结构命名不同例如静态名中不含encoder请自定义set_param_lr_fun否则所有参数 ratio 恒为 1.0。layerwise_decay的类型约束仅接受float或 PaddleVariable传入其他类型如numpy标量、int会抛出TypeError(coeff should be float or Tensor.)。权重衰减是解耦的weight_decay作用于参数本身param - param * coeff * lr而非梯度配合apply_decay_param_fun可精确控制哪些参数参与衰减例如常见的bias 与 LayerNorm 参数不做权重衰减策略。多精度训练建议fp16 训练时建议开启multi_precisionTrue此时衰减与更新都在 fp32 master weight 上进行数值稳定性更好同时配合lazy_mode可在超大参数下加速更新。API 文档定位本文对应的 API 文档页 paddlenlp.ops.optimizer.AdamwOptimizer.rst 为automodule自动生成页同系列可继续查阅 paddlenlp.ops.optimizer.adamwdl.rstAdamWDL 详细文档、paddlenlp.ops.optimizer.ema.rstEMA以及 paddlenlp/ops/optimizer/lr.pyInverseSquareRootSchedule反平方根学习率调度它们可组合出一套完整的 PaddleNLP 训练优化组件。九、总结AdamWDL是 PaddleNLP 提供的AdamW 分层学习率衰减一体化优化器它完整继承paddle.optimizer.AdamW的解耦权重衰减、多精度、梯度裁剪等能力并通过layerwise_lr_decay默认策略 可自定义的set_param_lr_fun回调在 Adam 算子执行前按层注入动态学习率。其实现中临时缩放学习率、用完即恢复decay_coeff 按学习率缓存、动转静每步清空等设计细节对理解 Paddle 优化器扩展机制亦有很强的参考价值。对于 BERT 系、GPT 系等深层 Transformer 的预训练与微调任务AdamWDL是一个开箱即用、参数语义清晰的选择。赞分享人工智能大模型预训练微调LoRARLHF强化学习分布式训练【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址https://gitcode.com/gh_mirrors/pa/PaddleNLP点击查看免费下载相关推荐PaddleNLP 优化器模块解析AdamWDL 分层学习率衰减、EMA 与逆平方根调度器实战指南PaddleNLP 优化器模块解析AdamWDL 分层学习率衰减、EMA 与逆平方根调度器实战指南 paddlenlp.ops.optimizer 是 Pad人工智能大模型预训练微调LoRARLHF强化学习分布式训练模型推理服务推理引擎模型量化模型压缩本地部署NLPPaddleNLP 中的 AdamWDL 优化器基于层间学习率衰减的 Transformer 训练方案PaddleNLP 中的 AdamWDL 优化器基于层间学习率衰减的 Transformer 训练方案 导读 AdamWDLAdamW with Dynam人工智能大模型预训练微调LoRARLHF强化学习分布式训练模型推理服务推理引擎模型量化模型压缩本地部署NLP如何使用 annotated_deep_learning_paper_implementations 配置 AdamW 权重衰减与 warmup 学习率调度如何使用 annotated_deep_learning_paper_implementations 配置 AdamW 权重衰减与 warmup 学习率调度 在人工智能深度学习大模型NLP计算机视觉强化学习LoRA创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考