这是我们过去9个月一直在努力的目标：让MoE训练速度提高约2倍，内存使用减少约2倍！亮点： - MoE通常在现代模型中占用最多的时间和内存。事实证明，可以通过数学方式重写MoE的反向传播，以减少在前向传播中需要存储的激活内存约2倍，从而在不额外进行矩阵乘法重计算的情况下获得相同的梯度。我非常喜欢这个结果，因为它结合了算法和系统的见解。 - 分析MoE层的瓶颈会导致一种自然的优化策略：尽可能减少内存的读写！收集前向传播的输入和反向传播的输出梯度有时会花费与分组GEMM相同的时间。我们将收集与分组GEMM融合，并重叠内存访问和计算，使整个层的速度提高约2倍。 - 计算专家路由的top-k可能会花费意外的时间，约占整个MoE层的15-20%！标准的top-k实现使用基数top-k算法，适合大k但对小k不够优化。我们使用比托尼克top-k算法重写了top-k，有时比pytorch的top-k快20-30倍！ 所有主要内核都是用Cute-DSL编写的，因此应该很容易扩展（并安装 :D）。Hopper内核已经发布，Blackwell内核也快准备好了。MoE模型的训练硬件效率曾经低于2倍，希望Sonic-MOE能改变这一点。