MindSpore 原生权重 EMA(指数移动平均)

📅 发布时间:2026/7/16 5:53:19 👁️ 浏览次数:
MindSpore 原生权重 EMA(指数移动平均)
大模型训练时权重会因批次噪声、梯度波动、动态批不断抖动直接导致验证精度忽高忽低模型容易过拟合最终推理效果不稳定。权重 EMA 是所有工业级大模型必开的单点稳定优化没有之一。它不修改训练流程、不增加计算量、不影响前向只在后台对权重做指数滑动平均用最平滑、最泛化的权重做验证与推理精度稳、泛化强、收敛曲线干净。核心原理一句话维护一份影子权重shadow weights每一步按衰减系数 shadowdecay×shadow(1−decay)×current 做指数平滑推理 / 验证时加载这份平滑权重彻底抹平训练抖动。MindSpore 原生极简代码import mindspore as ms import mindspore.nn as nn # 工业标准衰减率LLaMA/Qwen 通用 EMA_DECAY 0.999 class WeightEMA(nn.Cell): 大模型权重EMA纯原生API静态图完美支持 def __init__(self, network, decayEMA_DECAY): super().__init__() self.network network self.decay decay # 初始化影子权重和模型参数完全同形状 self.shadow_weights ms.Parameter( [w.clone() for w in network.trainable_params()], requires_gradFalse ) def update_ema(self): 核心一步完成EMA权重更新 for i, weight in enumerate(self.network.trainable_params()): self.shadow_weights[i] self.decay * self.shadow_weights[i] (1 - self.decay) * weight def load_ema_weights(self): 推理/验证时加载EMA平滑权重 for i, weight in enumerate(self.network.trainable_params()): weight.set_data(self.shadow_weights[i]) # 集成到训练只加一行零侵入 class EMATrainOneStepCell(nn.TrainOneStepCell): def __init__(self, network, optimizer): super().__init__(network, optimizer) self.ema WeightEMA(network) def construct(self, *inputs): loss super().construct(*inputs) # 每步自动更新EMA权重 self.ema.update_ema() return loss