MoE预训练中的专家相关放置与Token迁移

专家混合 (MoE) 层用 E 个专家网络替换 Transformer 的前馈块，并且每个Token都路由到这些专家中的 k 个。在专家并行 (EP) 下，专家分布在 GPU 上，每个 MoE 层在前向和后向传递中运行所有到所有的集合，将Token分派给专家，然后组合结果。在每个节点有 8 个 AMD Instinct MI300X GPU 的集群上，这些集合可以在 EP32 上使用 top-2 路由执行 45% 的训练步骤，使用 top-6 路由执行 60% 的训练步骤。我们发现，在预训练的早期，路由器已经学会了将Token分配给相关模式的专家，无论是在层内还是跨层。在 top-2 中，一层中 0.8% 的专家对被 42% 的 token 一起选择，并且 token 在一层选择的专家预测它在下一层选择的专家。我们使用这些相关性在Token自己的 GPU 上保留更多Token专家分配，这减少了 GPU 之间和节点之间的通信。相关的专家放置将经常被选择的专家放在同一个 GPU 上。与将每个Token发送到每个 GPU 一次的调度程序相结合，它可以删除多达 58% 的已调度行。当序列并行性在 EP 组中对Token进行分片时，会应用Token混洗。它将每个Token移动到 GPU，预计在注意力之后的减少分散期间容纳其下一层专家。在一个节点上，这将Token GPU 上服务的专家分配的份额从 12.5% 提高到 59%。在 Megatron-LM 中，从 8 到 64 的 EP 度，采用 top-2 和 top-6 路由，这两种方法将所有时间缩短了 1.16-2.63 倍，将端到端步骤时间缩短了 1.41 倍。这两种方法都不会改变模型的底层路由决策或专家参数。