FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models

A mixture-of-experts model saves computation by sending each token to only a few of its many expert networks, but the experts live on different machines, so the saving is paid back in network traffic — and training systems hide that traffic by overlapping it with computation; this paper (Hong Kong University of Science and Technology and collaborators) points out that existing systems overlap the wrong things, treating traffic inside a machine and traffic between machines as one resource and using a single pipeline setting for both directions of training, and shows that fixing all three oversights compounds to about 20 percent.

Why this problem is worth solving

The structure being trained is unusual enough that its costs need spelling out.

The second thing to establish is which of those patterns compete and which do not.

What people did before, and where it broke

• Overlap the all-to-all exchange with expert computation. The standard technique: split the batch, transfer one chunk while computing on another. Several systems do this and it removes a large part of the exposed communication.

• The gap. It considers only two things at a time. The transfers that gather a sharded expert within a node are left in sequence with the between-node exchange, even though they use a different link and could run alongside it.

• Overlap gradient synchronisation with the non-expert layers. Since gradients must be averaged across data-parallel replicas, systems start that transfer early and hide it behind attention and other dense computation.

• The gap. The expert layers are treated as unavailable, so the gradient traffic is packed into whatever dense work exists. But the expert schedule has gaps of its own, and gradients can be sliced to fit them.

• Use one pipeline setting for the whole step. The number of chunks the batch is split into is chosen once and applied to both forward and backward passes.

• The gap. Backward propagation does strictly more work, since it computes gradients with respect to weights as well as inputs, so its optimal chunk count differs. Across 1,458 configurations measured here, 912 wanted different settings for the two directions.

• Slice gradients into fixed-size chunks. One prior system does interleave gradient traffic more aggressively by cutting it into pieces of a preset size.

• The gap. A fixed size cannot match gaps that vary with layer shape and cluster, so the paper describes its performance as hit or miss, and measures it as slower than adaptive partitioning.

How a schedule gets chosen instead of assumed

The idea

Two kinds of communication are two resources, not one. The simplest and largest observation is that a mixture-of-experts step generates traffic on two separate fabrics. The all-to-all exchange that routes tokens to their experts crosses between nodes. Splitting a single expert across the accelerators inside one node produces gathering and scattering that never leaves the node, travelling on the much faster internal link. Existing schedules do not overlap these with each other, so the internal link sits idle while the network works and then the reverse, and the step is longer than the sum of its bottlenecks requires. Putting them on separate streams and pipelining them against each other is worth roughly five percent on its own in the ablation, which is modest until you notice it costs nothing but recognising the distinction. Decomposing a collective operation so it can be interleaved with the work that depends on it is the same manoeuvre as the compiler decomposition work covered elsewhere on this site, applied to two collectives against each other rather than to a collective against computation.

Forward and backward are different problems and deserve different answers. The pipelining knob is how many chunks the batch is split into, and every prior system picks one value for the whole step. But the backward pass computes gradients with respect to both inputs and weights, so it does more arithmetic per unit of communication than the forward pass, which shifts where the balance point sits. The paper measured this across 1,458 configurations and found 912 of them — nearly two thirds — wanted different chunk counts in the two directions. Solving separately is not conceptually deep, but it requires a cost model good enough to predict the difference, which is why the profiling step exists. Letting each part of a training step choose its own setting rather than inheriting a global one is exactly the argument of the adaptive pipeline work covered elsewhere on this site, which lets each stage recompute only as much as its own memory pressure demands.

Gradient traffic should be poured into the schedule, not scheduled around it. Data parallelism requires averaging gradients across replicas, which is between-node traffic and therefore contends directly with the all-to-all exchange. Existing systems respond by hiding gradient synchronisation behind the dense, non-expert parts of the model and treating the expert layers as off limits. The paper instead treats the gradient as a fluid: after the expert schedule has been chosen, whatever gaps remain are identified across all layers, and gradient slices are assigned to fill them, with a second optimisation distributing what is left over. This is what distinguishes it from a prior system that also interleaves gradient traffic but cuts it into fixed-size pieces — a fixed size cannot match gaps that vary with layer shape and cluster, which is why that approach is measured as slower here. Treating collective communication as something to be shaped to fit rather than accepted as given is the disposition behind the collective-language work covered elsewhere on this site.

Four schedules for the same layer

This is the baseline and nobody ships it, but it makes the structure visible. Six operations, of which two are within-node communication, two are between-node communication, one is computation, and one is the gradient synchronisation that data parallelism adds.

Why it works

• Different links can carry traffic at the same time. Within-node and between-node transfers use separate hardware → issuing them on separate streams lets both fabrics work at once → about five percent, obtained by recognising a distinction rather than by building anything.
• The two passes balance at different points. Backward propagation computes gradients for weights as well as inputs, so it does more arithmetic per byte → its optimal chunk count differs from the forward pass → 912 of 1,458 configurations wanted different values, which a single setting cannot provide.
• Gradients are arbitrarily divisible, so they make ideal filler. Slicing gradient synchronisation to match the gaps a schedule leaves → the between-node link stays busy without displacing the all-to-all exchange → this beats cutting gradients into fixed-size pieces, since gaps vary with layer shape and cluster.
• Simple cost models are accurate enough here. Fitting a startup cost plus a linear rate to each operation → predictions with correlation above 0.998 for computation and above 0.9999 for every collective → the scheduler can reason about alternatives without measuring them.
• The abstraction is what lets the scheduler be general. Defining six sub-modules with fixed interfaces and profiling them as black boxes → the scheduler needs no knowledge of how a routing function is implemented → four different routing functions and user-supplied ones are all schedulable by the same machinery.

What it costs (trade-offs and scope)

• Profiling is required per cluster. The cost models must be fitted with microbenchmarks before training. The paper measures this at under a hundred seconds of measurement plus milliseconds of fitting, done once, but it is a step that must happen.
• The scheduler assumes a particular alignment. The analysis presumes expert sharding matches the number of accelerators in a node, so within-node traffic really does stay within the node. That is a common configuration rather than a universal one.
• The gains are proportionally modest on the synthetic sweep. Across 1,458 configured layers the improvement is 18 to 22 percent. The much larger figures come from comparison against one baseline system on whole models, where the gap includes more than scheduling.
• The case analysis is hand-derived. Rather than optimising the schedule directly, the paper enumerates four regimes by which resource dominates and solves each. This is tractable but it is a taxonomy someone constructed, not a general solver.
• Modest clusters. The evaluation runs on two testbeds of up to 48 accelerators, with consumer and workstation-class cards. Behaviour at the scale where these models are actually trained is not shown.

The evaluation

Check yourself

Sources

FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models. Xinglin Pan, Wenxiang Lin, Lin Zhang, Shaohuai Shi, Zhenheng Tang, Rui Wang, Bo Li, Xiaowen Chu. ASPLOS 2025, Volume 1, pages 524 to 539. DOI 10.1145/3669940.3707272. Full text read from the authors’ preprint, arXiv 2501.10714.