Stable-MoE: Lyapunov-based Token Routing for Distributed Mixture-of-Experts Training over Edge Networks

TL;DR

This paper introduces Stable-MoE, a Lyapunov-based token routing framework for distributed mixture-of-experts training on edge networks, improving efficiency, stability, and performance amid resource heterogeneity and stochastic workloads.

Contribution

It presents a novel Lyapunov optimization approach for real-time token routing and resource allocation in edge MoE systems, ensuring stability and enhanced throughput.

Findings

Achieves at least 40% higher system throughput.

Improves test accuracy by at least 5%.

Demonstrates stability and efficiency in resource-heterogeneous environments.

Abstract

The sparse activation mechanism of mixture of experts (MoE) model empowers edge intelligence with enhanced training efficiency and reduced computational resource consumption. However, traditional token routing in distributed MoE training faces significant challenges in resource-constrained edge networks characterized by heterogeneous computing capabilities and stochastic token arrivals, which inevitably suffer from workload backlog, resource inefficiency, and performance degradation. To address this issue, we propose a novel Lyapunov-based token routing framework for distributed MoE training over resource-heterogeneous edge networks, termed Stable-MoE. Specifically, we formulate a stochastic optimization problem to maximize both system throughput and gating consistency via optimizing the token routing strategy and computational resource allocation, while ensuring long-term stability of both token and energy queues at the edge devices. Using the Lyapunov optimization, we transform the intractable long-term optimization problem into tractable per-slot subproblems by enabling online decision-making of token routing and computation frequency utilization without the knowledge of future system states. Experimental results on the SVHN and CIFAR-100 datasets demonstrate that Stable-MoE outperforms the baselines with at least 40% and 5% gains in system throughput and test accuracy, respectively.