Sonar-MoE：通过反向传播期间的专家更新实现内存高效的混合专家训练
Sonar-MoE: In-Backward Expert Updates for Memory-Efficient Mixture-of-Experts Training
精选理由
Sonar-MoE在保持训练效果的同时降低峰值显存、缩短训练时间，让大规模 MoE 微调更高效。
AI 摘要
混合专家（MoE）模型的大多数参数属于专家，而标准训练会将所有专家梯度保留在 GPU 内存中，直到优化器执行更新。这部分内存占用随专家数量和规模的增加而增长，因此会随着 MoE 模型的扩展而上升。我们提出 Sonar-MoE，它在反向传播期间更新每个专家张量，并立即释放其梯度。该方法将稠密模型中的反向传播期间优化器更新机制扩展到批量化的 MoE 专家张量，使 GPU 中同时保留的专家梯度始终不超过一层。我们使用两种专家计算内核后端，以 bf16 精度微调 OLMoE-1B-7B 和 Granite-3.1-3B-A800M 时，Sonar-MoE 将峰值内存占用降低了 19% 至 22%。在不进行专家梯度裁剪的情况下，其更新结果与标准训练逐位一致，且留出集损失与所有基线的差值均不超过 0.003。与采用全局范数裁剪的标准训练相比，每步耗时也减少了 17% 至 23%。更新整个批量张量而非逐个更新专家，是 Sonar-MoE 保持高速度的关键：它在内存占用相同的情况下，避免了 20% 至 72% 的速度下降。Sonar-MoE 在内存占用和速度上与对所有参数采用反向传播期间优化器更新的方法相当，同时仍可对稠密参数进行梯度裁剪。