TAOT: No More Waiting on "the Slowest GPU" in Large-Scale MoE Training

#ai

The long-tail problem in MoE training

MoE is already the mainstream form of large models. Each token activates only a few experts, so model size can keep scaling up without per-step computation growing along with it. But expert parallelism (EP) has a practical problem: how many tokens each GPU has to compute is decided on the fly, entirely by the routing result. Once a few experts become hot spots, tokens converge on the ranks hosting them, the overall time of synchronous training is held back by those few GPUs, and the rest can only wait idle after finishing.

What existing methods overlook

The industry thought of a fix for this early on: if one expert has too much to compute, make a temporary copy of its weights, place it on an idle GPU, and let that GPU take on part of the work.

The problem is that these methods optimize essentially one thing — how evenly the load is spread — and assume that among the idle GPUs it does not matter which one you pick. But real machines are not like that:

GPUs within the same machine talk to each other over a high-speed channel like NVLink, so moving data is fast and cheap;
GPUs across machines have to go over the network (InfiniBand), and moving the same amount of data is far slower and far more expensive.

This brings a consequence that is easy to overlook: even if two schemes spread the load equally evenly, as long as one keeps the weights on the local machine and the other moves them to a different machine, the communication overhead actually spent can differ by a wide margin.

The TAOT idea: count the cost of moving, too

Baidu's LoongForge team argues that which GPU a replica ends up on cannot be a matter of asking only "where is there room" — you also have to ask "how expensive is it to move there".

This is exactly where TAOT starts. It is the first to write both the peak-shaving gain and the cost of moving weights across nodes into the optimization objective, meaning that when picking a landing spot for a replica it weighs two things at once: on one side, flatten the load; on the other, spend as little cross-machine communication cost as possible.

It borrows the idea of "optimal transport" from mathematics to do this. You can think of it as a scheduler that keeps track of the communication bill:

Stay local if you can: if there is still room on the local machine, use it first and save the trip outward;
Go outward only when there is really no room left: once the local machine is full, crossing machines is not entirely forbidden either — in its view this is simply an option that is "available, but pricier", rather than being ruled out across the board.

Existing methods land at roughly two extremes. One simply ignores topology and puts a replica on whichever GPU has room, so replicas easily get thrown onto the machine next door: the load is flattened, but cross-machine communication climbs sharply. The other goes the opposite way and hard-codes where weights are allowed to move with a fixed topology graph; at scale, hot spots get broken up, nearby capacity is used up first, and the distant GPUs that are in fact still idle are kept out by the graph, so the schedulable space actually narrows.

TAOT takes the middle road: it does not forbid crossing machines, nor does it hard-code the paths. It simply puts a "pricier" tag on going cross-machine, so that it naturally ranks lower when the bill is added up. Hence the larger the scale and the more idle GPUs there are, the more room it has to maneuver. Below is the overall system architecture.

The overall TAOT flow: once routing is computed, the load of each GPU is collected first and handed to the three-phase planner, which settles on a replica plan; at execution time the replica weights are moved into place, and that transfer is run in parallel with computation on the GPU.

Solving the two objectives, "balance" and "communication cost", together makes the problem size grow rapidly with EP degree. TAOT breaks it into three steps, coarse to fine, each answering just one question:

Phase 1 (rank level): roughly where should the load flow? Take overload as supply and spare capacity as demand, add the communication-cost matrix, and solve for optimal transport, which yields a global flow suggestion table. It is only a soft hint; its role is to give the next two steps a global reference so they do not just go by what is in front of them.
Phase 2 (expert level): which expert's replica goes into which GPU? Turn the continuous flow into a 0/1 decision: for every "(idle GPU, hot expert)" pair, compute a score made up of three weighted terms — peak-shaving gain, how far away the landing spot is, and the Phase 1 flow hint — and take the highest-scoring combination.
Phase 3 (token level): which GPU sends which of these tokens? The tokens of one expert are scattered across multiple GPUs, so round after round of bidding settles how many each source GPU sends. The winner's price goes up, so it naturally yields in the next round, and locality is thereby carried into the assignment.

A good algorithm alone is not enough; the engineering has to keep up

For a method to really run inside training, a good algorithm alone is not enough; two more things on the engineering side have to keep up.

First, the scheduling itself has to be fast enough. This planning is not computed once and done with; it has to run live, over and over, one microbatch after another throughout training. If the planning is itself slow, the time saved on the communication side gets spent on scheduling instead. To that end the team did extensive optimization at the kernel level, compressing what originally took hundreds of fragmented computations down to a few, and finally brought the overhead of the whole planning process to within 1% of forward computation time.

Second, "hide" the act of moving inside the computation. Replica weight distribution is placed on the communication stream, in parallel with the GEMM of the home experts on the same GPU; on the backward pass, guest gradients are returned and accumulated through a single All-to-All. The transfer time is largely covered by computation.

How well does it actually work

Faster: on Qwen3-30B-A3B, forward-and-backward computation time per iteration drops from 155.4ms to 108.8ms, a 1.43× end-to-end speedup;
Cheaper: given equally even load, its communication cost is the lowest among all the SOTA methods compared, up to 74% below the competition;
No accuracy sacrificed: this speedup comes from more sensible scheduling, not from lowering numerical precision; the loss curve nearly coincides with that of the standard scheme;
More pronounced at larger scale: the larger the parallel scale and the more imbalanced the initial load, the more significant the speedup TAOT brings. From EP4 to EP16, the speedup reaches up to 1.79×; as initial imbalance rises from 30% to 90%, the speedup grows from 1.21× to 1.75×. The harder the scenario, the greater the gain.

Closing: a scheduler that keeps track of the communication bill

Dynamic replicas are already a fairly mature line of work in MoE balancing, with Echo, LPLB, and LLEP all on it. What sets TAOT apart is the optimization objective: it weighs the communication cost of a replica's landing spot against the peak-shaving gain within one and the same objective function, instead of only seeking an even load. The implementation uses a continuous cost matrix together with a soft topology preference, writing "prefer intra-node, cross node when necessary" directly into the objective — no reliance on special hardware, no confinement to a single node, no predefined graph.

If you are doing large-scale MoE training and your training is held back in the long tail by a handful of hot experts, TAOT is worth a try. It has been integrated, as an extension to the MoE expert mechanism, into LoongForge, the omni-modal training framework open-sourced by Baidu Baige.