Papers

Sonar-MoE:在反向传播中更新专家参数,降低 MoE 训练显存

Sonar-MoE: In-Backward Expert Updates for Memory-Efficient Mixture-of-Experts Training

Why we picked it

论文讲了个省显存的训练技巧:专家参数反向传播时就地更新,微调 OLMoE 这类 MoE 模型显存省约两成,速度还更快。

Sonar-MoE 将 optimizer-in-backward 扩展到 MoE 的批量专家张量上,在 backward 阶段就地更新每个专家张量并立即释放梯度,GPU 峰值显存最多只需保存一层的专家梯度。在 bf16 微调 OLMoE-1B-7B 和 Granite-3.1-3B-A800M 的实验中,峰值显存降低 19% 到 22%,每步耗时比标准训练加 global-norm 裁剪快 17% 到 23%。不裁剪专家梯度时,更新结果与标准训练逐位一致,held-out loss 与基线差距在 0.003 以内。相比逐个专家更新的做法(会慢 20% 到 72%),它通过整批张量更新保住了速度。

Original · arXiv cs.LG

Sonar-MoE: In-Backward Expert Updates for Memory-Efficient Mixture-of-Experts Training

Most parameters of a Mixture-of-Experts (MoE) model belong to its experts, and standard training keeps every expert gradient in GPU memory until the optimizer step. This memory grows with the number and size of the experts, so it rises as MoE models scale. We introduce Sonar-MoE, which updates each expert tensor during the backward pass and frees its gradient right away. It brings optimizer-in-backward from dense models to batched MoE expert tensors, so the GPU never holds more than one layer of expert gradients. When we fine-tune OLMoE-1B-7B and Granite-3.1-3B-A800M in bf16 with two expert-kernel backends, Sonar-MoE lowers peak memory by 19% to 22%. Its updates are bitwise identical to standard training without expert gradient clipping, and held-out loss stays within 0.003 of all baselines. Each step also takes 17% to 23% less time than standard training with global-norm clipping. Updating whole batched tensors rather than individual experts is what keeps Sonar-MoE fast, as it avoids a 20% to 72% slowdown while using the same memory. Sonar-MoE matches optimizer-in-backward over all parameters in memory and speed while still clipping the dense parameters.