ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

昇思MindSpore大模型优化器参数详解:从Adam到分布式训练

昇思MindSpore大模型优化器参数详解:从Adam到分布式训练 昇思 MindSpore 大模型优化器参数详解打开一份典型的大模型训练脚本大部分人的注意力都放在模型结构、数据集、分布式并行配置上优化器那一块往往是从某个开源项目里复制过来的几行参数。我过去也这样干过直到有次用 MindSpore 训练一个十亿级参数的生成模型loss 在 1.8 附近死活降不下去改学习率、换调度策略、调整批大小都没用最后逐条排查发现是 Adam 的 beta2 用了默认的 0.999导致二阶动量估计对近期梯度变化太迟钝整个优化器陷入了一种“看起来在更新、实际上步子已经迈不开”的状态。从那时起我就明白在大模型训练里优化器参数不是填个表格交差的事它直接决定你的训练曲线是平缓收敛、剧烈抖动还是直接飞掉。这篇内容我就以 MindSpore 为背景把大模型训练中优化器选型、核心参数、动态调度、分布式状态这几个关键环节掰开揉碎讲一遍里面也会带上我在实际训练里踩过的坑和一些调参经验希望对正在折腾大模型的朋友有帮助。1. 先弄清优化器在大模型中扮演的角色为什么它经常被低估1.1 大模型场景下优化器比很多人想象的更“重”模型规模大了之后优化器不再只是“把梯度方向算出来然后更新权重”这么简单。首先优化器参数的数量会直接决定训练状态占据的显存规模。以 Adam 为例每个参与训练的参数都要额外保存一阶动量m和二阶动量v如果优化器状态用 FP32 存储仅这两项就是模型参数量字节数的 8 倍。一个 7B 参数的模型光优化器状态就要占 224GB 的显存这还没算模型本身和梯度。所以在大模型训练里优化器的选型和参数配置首先就要回答“我能不能放得下”的问题。其次优化器的数值行为会影响整个训练的稳定性。大模型的 loss landscape 往往非常陡峭且存在大量“坏区域”学习率稍微大一点就可能撞进 NaNbeta 参数不合适又可能导致收敛极慢。换句话说大模型训练对优化器超参数的要求比中小模型严格得多一个小参数的偏差会被放大到整个训练过程中。1.2 MindSpore 优化器选型图谱SGD、Adam、AdamWeightDecay 与 LAMB 的适用边界MindSpore 的mindspore.nn里提供了一系列优化器我在实际项目里主要用过四类SGD、Momentum、Adam以及它的变体AdamWeightDecay、还有面向超大批量训练的LAMB。SGD / Momentum在小规模任务或者迁移学习微调浅层网络时还能用但在大模型预训练场景下基本不推荐。SGD 对学习率的敏感度极高且没有逐参数的自适应调整能力在稀疏梯度占比较高的 Transformer 结构里会出现明显的收敛迟缓。Adam经典的逐参数自适应方法利用一阶动量估计梯度方向、二阶动量估计梯度尺度的缩放在 NLP、CV 的各种任务上表现都相当稳健。MindSpore 的nn.Adam已经做了针对昇腾硬件的算子融合优化实测训练吞吐比逐个算子组合高出不少。AdamWeightDecayAdamW这是大模型预训练最常用的默认优化器。它与 Adam 的核心区别在于把权重衰减从 L2 正则中解耦出来不参与动量计算而是直接在参数更新时乘一个衰减系数。近几年的工作流里基于 Transformer 的大模型尤其 GPT 风格模型基本都会选择 AdamW。LAMB在超大批量batch size 从几千到几万的条件下LAMB 能以较大的学习率稳定训练因为它会对每一层参数单独归一化更新幅度。做超大规模预训练时可以考虑但它的敏感参数比较多需要更细致的调参不能无脑套用。选优化器不是“哪个新选哪个”核心要看三件事你的模型结构是否有稀疏/长尾参数、你的批量大小是否已经大到影响梯度噪声水平、你的显存是否能装下对应优化器的状态。2. 五个核心参数逐一拆解learning_rate、beta、weight_decay、eps 的作用边界2.1 learning_rate一个看似简单但最容易被误用的参数学习率是整个优化器里最显式的一个旋钮但它在大模型训练里并不是“越大越快、越小越稳”。我习惯把大模型训练的学习率分成几个档位看预训练从头开始常见在1e-4到3e-4之间具体取决于批量大小和数据质量。批量越大梯度估计越稳定学习率可以适当上调。领域微调或指令微调通常在1e-5到5e-5之间因为预训练权重已经进入了较好的局部区域学习率过大会直接破坏已有参数结构。LoRA 这类参数高效微调基础学习率可以稍大约1e-4到2e-4但需要配合适当的层分配策略。有一个实际经验当你在小规模试跑时发现 loss 在某个值附近反复震荡、不下降也不发散很多人第一反应是加大学习率实际上大模型场景里更可能是学习率偏大导致在最优解附近来回“弹跳”。把学习率降到原来的 1/5 到 1/10往往比换模型结构更有效。MindSpore 中可以这样设置固定学习率import mindspore as ms from mindspore import nn optimizer nn.AdamWeightDecay( paramsnet.trainable_params(), learning_rate2e-5, weight_decay0.01 )但要注意实际大模型训练几乎不会用固定学习率会配合动态调度这部分在下一节详述。2.2 beta1 与 beta2一阶、二阶动量的衰减策略决定了优化器的“记忆长度”Adam 的更新公式里beta1控制一阶动量梯度均值的指数衰减速率beta2控制二阶动量梯度平方均值的指数衰减速率。用直白的话解释beta10.9意味着优化器大约会参考最近 10 步的梯度方向beta20.999则意味着二阶动量会考虑近 1000 步的梯度幅度信息。大模型场景里beta10.9基本是共识很少动它。而beta2是一个非常容易被忽略但影响极大的参数。当beta20.999时二阶动量估计非常平滑适合数据分布相对稳定的场景。当数据分布发生变化较快比如训练新领域数据、长文本、对话类数据混合或者出现较大的梯度波动时过大的beta2会让二阶动量“反应迟钝”导致优化器无法快速调整步长loss 看起来就卡住了。我实际踩过的例子是用 MindSpore 训练一个多任务混合数据的大模型loss 在过了某个阶段后持续横盘后来把beta2从0.999调成0.95同样的总步数下收敛速度明显提升。原因是混合数据带来的梯度方差更大过快平滑的二阶动量会把有用的梯度信号“抹平”。MindSpore 中写法optimizer nn.AdamWeightDecay( paramsnet.trainable_params(), learning_rate2e-5, beta10.9, beta20.95, eps1e-6, weight_decay0.01 )2.3 weight_decayAdam 和 AdamW 的关键区别以及哪些层不该衰减很多人把weight_decay等同于 L2 正则这在 SGD 里差不多但在 Adam 里有一个经典问题L2 正则的梯度会先被加到参数梯度上再参与一阶、二阶动量的计算这使得权重衰减的实际效果会被自适应学习率“稀释”尤其是参数本身尺度较小的时候。AdamWeightDecayAdamW的做法是解耦权重衰减更新参数时直接从当前参数中减去一个与梯度无关的weight_decay * param项。这样权重衰减就能真正发挥抑制过拟合的作用。在大模型实践中weight_decay的常用值在0.01到0.1之间。但有一个细节容易被忽略并不是所有参数都需要做权重衰减。归一化层的 gamma/beta、偏置项这些参数本身不参与权重规模的累积如果对它们也做 decay反而可能引入不必要的限制甚至导致训练不稳定。MindSpore 的AdamWeightDecay支持通过decay_filter来精细化控制哪些参数执行衰减def decay_filter(param): if layernorm in param.name.lower() or bias in param.name.lower(): return False return True optimizer nn.AdamWeightDecay( paramsnet.trainable_params(), learning_rate2e-5, beta10.9, beta20.95, eps1e-6, weight_decay0.01, decay_filterdecay_filter )这个细节在中小模型上可能感觉不明显但在大模型上千亿参数规模的 LayerNorm 和 bias 参数量并不少正确的衰减过滤会把最终效果和稳定性都提升一个档次。2.4 eps数值稳定性与低精度训练的平衡点eps是加法到二阶动量上的极小常数作用就是防止除零。默认值1e-6在标准 FP32 训练下够用但在混合精度训练FP16 或 BF16成为大模型标配的今天eps的选择就值得重新审视。FP16 能表示的数值范围有限如果eps设得太小比如1e-8二阶动量在某些梯度很小的参数上会等于 0导致更新量变为极不稳定的除零结果。反过来如果eps设得太大比如1e-4又会压缩自适应学习率的动态范围让优化器偏向于“伪 SGD”的行为影响收敛精细度。我的经验范围是FP32 训练用1e-6FP16 混合精度训练如果梯度裁剪和 loss scaling 已经正常工作1e-6也可以但训练初期如果出现loss NaN优先检查eps是否被调低到了1e-8以下BF16 因为指数位更宽eps的敏感度会低一些但仍然不建议设到1e-8以下。3. 学习率调度与批量大小的配合不是所有“5e-5”都长一个样3.1 预热warmup为什么是必修课刚开始训练时模型权重还是随机的梯度方向包含大量噪声。如果一开始就用较大的学习率优化器会被带进一个比较差的参数区域后面再想拉回来非常困难。预热阶段用很小的学习率“让模型先站稳”再逐步增大到目标学习率这相当于给一个刚睡醒的人慢慢睁开眼看路而不是直接拉到强光下。MindSpore 生成动态学习率的方式比较直接可以用nn.dynamic_lr下的系列函数构造学习率列表再传给优化器。一个典型设置import mindspore as ms import mindspore.nn as nn total_steps 20000 warmup_steps 1000 base_lr 2e-5 lr_schedule nn.dynamic_lr.cosine_decay_lr( min_lr2e-6, max_lrbase_lr, total_steptotal_steps, step_per_epoch1, decay_epochtotal_steps, warmup_stepwarmup_steps ) optimizer nn.AdamWeightDecay( paramsnet.trainable_params(), learning_ratelr_schedule, beta10.9, beta20.95, eps1e-6, weight_decay0.01 )这里有几个点需要特别留意warmup_step必须和实际训练步数对齐不要估算。如果训练提前终止或中途改变数据 epoch 数学习率序列会失配。min_lr一般取base_lr的 1/10 左右不要设成 0。完全衰减到 0 在后半程会让优化器丧失微调能力最后几轮训练基本在“原地踏步”。3.2 余弦衰减与线性衰减的真实差异大模型训练后期到底应该用余弦衰减还是线性衰减业界没有绝对定论但经验上有一些倾向余弦衰减会让学习率在后半段平滑降低模型有机会在局部区域精细搜索适合训练步数比较充足、数据质量比较高的场景。线性衰减逻辑更“硬核”降到某一步直接归零适合训练预算紧张、希望在后段快速收敛的情况。我在 MindSpore 里尝试过两种方式结合实验对比预训练阶段我更倾向余弦衰减它的平滑特性在长时间训练里更容易保持稳定性而在领域微调或指令微调阶段因为总步数本身不多线性衰减在实践中也够用。关键原则是学习率衰减的速率不能过快否则优化器还剩大量“动量惯性”时学习率已经归零会导致最终效果明显变差。3.3 梯度累积改变的是什么当单卡显存放不下足够大的 batch 时通常会使用梯度累积把多个小 batch 的梯度累加后再更新一次参数。在优化器眼里这等价于把有效 batch size 变大了若干倍。有效 batch size 变大后梯度估计更平滑学习率的上限理论上是可以提高的。但需要注意Batch size 增大时学习率并非必须线性增长。在超大批量场景下优化器的自适应机制会弱化这种线性关系的适用性盲目随 batch 翻倍加大学习率很容易导致训练震荡。我的做法是batch size 翻一倍时学习率先加 30%-50%观察 loss 曲线的抖动情况如果训练不稳退回原学习率而不是强行追高。4. 分布式训练下优化器状态显存占用、混合精度与并行策略4.1 大模型训练中优化器状态究竟占了多少显存很多人一开始以为大模型显存占用“大头是模型参数”实际训练几次就会明白优化器状态才是那个最占地方的隐形大户。以 FP32 的 Adam 为例每个参数需要保存一阶动量4 字节、二阶动量4 字节加上参数本身4 字节和梯度4 字节一个参数就要吃掉 16 字节。如果模型 7B 参数单卡训练需要的裸数据量就是 112GB这在单卡环境下根本不可能。好在大模型训练几乎都会走分布式并行。MindSpore 的数据并行、模型并行都支持将优化器状态切片将不同分片放在不同设备上这样每个设备只要维护自己负责的那部分参数对应的优化器状态显存压力大幅下降。专业一点说这就是 ZeRO 优化器状态分片思路。配置这类并行时建议先统计出模型参数总量再反推每个设备的可用显存最后决定并行切分的维度。4.2 混合精度下的优化器参数FP16 梯度与 FP32 状态的配合混合精度训练里通常模型参数和梯度用 FP16/BF16但优化器状态会保留 FP32 副本。为什么要保留 FP32因为更新步骤中如果直接在 FP16 上做动量累积精度损失会随着迭代不断累积最终参数更新方向都会被误差淹没。MindSpore 中使用amp模块可以方便地把模型和优化器组装起来from mindspore import amp train_network amp.build_train_network( networknet, optimizeroptimizer, loss_fnloss_fn, levelO2, loss_scaleramp.DynamicLossScaler(scale_value1024, scale_factor2, scale_window2000) )在这种配置下eps和 loss scaling 是有联动关系的。梯度经过 loss scaling 放大后再反传相当于优化器看到的梯度幅度整体被放大此时如果eps相对梯度尺度太小二阶动量的数值稳定性指标会完全被 loss scaling 支配导致极端情况下出现异常。我的经验是开启动态 loss scaler 时优先让eps保持1e-6级别不要在低精度训练里把eps随意改小。4.3 优化器算子融合与状态缓存吞吐量优化的一个易忽略点MindSpore 在昇腾硬件上做了大量优化器级算子融合理论上你只需要调用nn.AdamWeightDecay这类高层 API底层会自动选择融合策略。但有一个经验仍值得分享当网络里存在大量分组参数时把参数切得过于零碎会降低优化器算子融合的效率因为每个参数组都要走一次单独的更新内核。如果你发现自己训练吞吐上不去且模型结构里有大量重复模块可以考虑对参数分组做合并处理或者在定义网络时考虑共享参数/统一命名空间让优化器可以更高效地处理参数批次。这是很多文档不会提的“隐形优化点”实测在相同硬件上能把吞吐提升 10%-15%。5. 优化器引起的训练事故排查从 loss 不降到训练崩坏的实战记录5.1 训练初期 loss 直接发散成 NaN 的排查链路如果模型结构本身没有问题可以先用几层全连接网络跑同数据验证那 loss 发散大概率出在优化器参数组合上。我的排查顺序通常是看学习率是不是过大。尤其从开源项目复制来的参数不见得适配你的模型规模和批量大小。先把学习率降到原来的 1/10 试跑 100 步。看 warmup 是否缺位。随机初始化权重加上大学习率往往第一波更新就把参数推到极端区域。看 eps 是否过小。在混合精度训练中eps1e-8很容易成为 NaN 的温床。看梯度裁剪配置。大模型里梯度范数裁剪建议设置为1.0左右在 MindSpore 中可以用nn.GradNormClip这类封装来限制极端梯度。下表是我实际遇到过的几种问题与对应解决方向现象常见原因优先调整项第 1-2 步 loss 直接变成 NaN学习率过大 / 缺少预热降低学习率增加 warmup 步数前几步正常几百步后突然发散二阶动量估计受极端梯度污染增大 beta2如 0.999调整梯度裁剪loss 在某个平台期横盘不动beta2 过大导致动量反应迟钝将 beta2 从 0.999 下调到 0.95 左右混合精度下持续小幅度抖动eps 与 loss scaling 不匹配保持 eps 在 1e-6 级别检查动态 loss scaler5.2 训练后期收敛过慢优化器侧的修正经验训练到后期loss 下降趋缓是正常现象但如果“缓”得太离谱就要从优化器参数里找原因。最常见的是学习率衰减到了极小值后优化器仍然在按固定动量方向更新造成参数在最优解附近“绕圈子”。这时可以检查学习率是否真的按 schedule 降到位以及weight_decay是否导致参数被过渡压缩。另一个容易忽略的点是不要频繁尝试在训练中途修改优化器超参数并保存 checkpoint。优化器的动量状态是从零开始累积的你改了 beta1、beta2相当于让一个已经养成运动习惯的物体突然换了一双新鞋前期动量记忆全部作废短期会有一段剧烈调整期。如果非要改建议带着优化器状态一起加载而不是只加载模型权重后重新初始化优化器。5.3 长期训练的“体检”优化器状态也要进 checkpoint训练大模型动不动就是几周甚至几个月中间机器故障、断点是家常便饭。很多人的 checkpoint 只保存模型权重结果恢复训练后发现 loss 曲线比之前差了一截。原因很简单优化器的动量与二阶动量没有被保存重新初始化的优化器在恢复训练后需要一个很长的“热身期”才能重建动量。所以在 MindSpore 里保存 checkpoint 时我会同时把优化器状态写进去import mindspore as ms ckpt [ {name: model, data: net.parameters_dict()}, {name: optimizer, data: optimizer.parameters_dict()}, {name: step, data: current_step} ] ms.save_checkpoint(ckpt, train_ckpt.ckpt)恢复训练时先把模型参数和优化器参数一起加载再继续跑训练曲线就能平滑衔接。这个细节在长周期训练中非常关键我建议从一开始就把优化器状态纳入存档机制而不是事后补。优化器参数看起来是训练脚本里的几行代码但它贯穿了显存规划、训练稳定性、收敛速度和故障恢复全过程。我在 MindSpore 上训练大模型的经验是先明确优化器和核心参数的范围再把学习率调度、混合精度、分布式分片这些环节联动起来看最后把所有超参和 checkpoint 策略固化到训练框架里而不是每次重启后靠感觉重新调。这样即使训练过程中出现异常你也能第一时间定位到问题在优化器还是在模型还是在数据上。
返回列表