开源项目

Megatron修复GatedDeltaNet在TP训练中的归一化梯度分歧

Sum GDN out_norm gradients across TP ranks without sequence parallelism

AI 基础设施模型架构SSM分布式训练基础设施

概述

GatedDeltaNet的输出归一化权重被各个值头共享,而张量并行把头拆到不同rank。旧代码只在序列并行打开时汇总相关梯度;TP大于1且SP关闭时,每卡使用局部头产生的部分梯度更新同一参数副本,造成静默分离。本次为GDN及GDN2的RMSNorm和LayerNorm参数增加跨TP求和标识,纳入既有SUM桶,并避免同时有SP标识时重复归约。来源给出两卡复现与对照,修复后梯度符合跨头求和参照。它针对明确的训练正确性边界,不是新模型发布或只在数值容差内的普通算子提速。