Spindle: Efficient Distributed Training of Multi-Task Large Models via Wavefront Scheduling

Yujie Wang, Shenhan Zhu, Fangcheng Fu, Xupeng Miao, Jie Zhang, Juan Zhu and 3 others

A modern foundation model handles text, images, audio and more, with separate encoders per modality feeding shared layers, and training it on a GPU cluster with tools built for one uniform model leaves GPUs idle because the pieces differ wildly in size and must wait on each other — Spindle breaks a training step into waves, decides how many GPUs each piece of the model deserves from measured scaling curves, packs the pieces into waves so every GPU is busy, places them to keep traffic inside NVLink islands, and trains up to 71 percent faster than Megatron-LM and DeepSpeed on models with up to ten tasks.

The argument in one glance

Problem

One model, many unequal parts, and they wait on each other

A multi-task multi-modal model such as ImageBind, OFASys or Qwen-VL has an encoder per modality (text, vision, audio, depth, thermal, motion) and a shared cross-modal module. Each training task activates a different subset. The parts differ in cost by an order of magnitude and in how well they speed up with more GPUs, and the shared layers cannot start until every encoder feeding them has finished. Training systems built for one uniform model see none of this.
What people did

Run tasks sequentially on the whole cluster, or split the cluster by task

Megatron-LM and DeepSpeed give the entire cluster to one task at a time. Task-level schedulers such as an Optimus-style allocator hand each task a fixed share of GPUs by marginal gain. DistMM, built for single-task multi-modal models, sizes each encoder within a task and runs tasks one after another.
Why it fell short

Whole-cluster execution starves light parts; task-level splits ignore the parts inside a task

Parallelizing a small audio encoder across 16 GPUs leaves its kernels nearly idle; the paper measures the resulting utilization swinging between tasks and within tasks over a step. Task-level allocation fixes only the between-task half, and as light tasks finish early their GPUs sit idle for the rest of the step. DistMM fixes only the within-task half and still serializes tasks. None reaches a global plan.
This paper's idea

Contract the graph, fit scaling curves, solve allocation per level, pack into waves, place for locality

Spindle contracts runs of identical operators into MetaOps and groups them by dependency depth into MetaLevels. It profiles each MetaOp for five minutes and fits a piecewise curve of execution time against GPU count. For each MetaLevel it solves a malleable project scheduling problem whose continuous optimum has every MetaOp start and finish together, found by bisection, then rounds each allocation to two valid integer GPU counts. A wavefront scheduler greedily packs slices of MetaOps into waves that fill all GPUs with aligned durations. A placement step keeps heavy data flows inside NVLink islands and balances memory. On an A800 cluster, Spindle trains Multitask-CLIP, OFASys and a 9.25-billion-parameter Qwen-VAL up to 71 percent faster than Megatron-LM and DeepSpeed, within 7 percent of the theoretical optimum, with plans computed in under three seconds.

Why this problem is worth solving

Meta’s Llama 3 is a multi-modal model with over 400 billion parameters trained on more than 48,000 GPUs. At that scale, every percent of idle GPU time is a large sum of money. The paper’s opening measurement shows how much idleness the model structure itself creates. Figure 1 traces GPU utilization while four tasks of a multi-modal model run one after another on a cluster: throughput swings from over 100 TFLOP/s to near zero within each task, because a task’s vision encoder is heavy and its text encoder is light, and swings again between tasks, because a motion task and a text task have nothing like the same cost.

The cause is

workload heterogeneity

Differences in computational cost and in how well cost falls with more GPUs, both between tasks and between the operators inside one task.

combined with

execution dependency

The requirement that an operator cannot start until every operator whose output it consumes has finished.

: shared layers must wait for every encoder that feeds them. A training system that gives the whole cluster to one component at a time is fast for the heavy ones and wasteful for the light ones. The obvious repair, running components side by side on separate GPUs, raises the questions Spindle answers: how many GPUs each deserves, in what order, on which devices, and how to keep shared parameters consistent.

What people did before, and where it broke

Megatron-LM and DeepSpeed. Mature 3D-parallel training for a single model. Applied to a multi-task model, they run each task’s sub-model across the whole cluster in sequence. Where it broke. Light operators are spread over all devices. The paper’s example: DeepSpeed runs a small audio operator on 16 GPUs and leaves the kernels underused, where four would do.
Task-level allocation (Spindle-Optimus). Adapt Optimus’s marginal-gain rule to give each task a share of GPUs. Where it broke. It sees tasks, not the operators inside them, so a task’s heavy and light parts get the same share, and when short tasks finish their GPUs idle for the rest of the step. Up to 44 percent faster than DeepSpeed, but well short of the optimum.
DistMM extended to multi-task (DistMM-MT). Allocate resources to the encoders within one task, then run tasks sequentially. Where it broke. Single-task by design; on OFASys, whose text adaptor is light, its intra-task parallelization gains almost nothing, and it never optimizes across tasks. Up to 20 percent faster than the SOTA systems.
Automatic parallelizers: Alpa, Unity, Galvatron. Search for hybrid parallel plans. Where it broke. They assume a spatial pipeline over a homogeneous model, with every operator of a stage on the same devices.
Cluster schedulers: Optimus, Pollux, Gandiva. Allocate GPUs across independent jobs by scalability. Where it broke. Jobs are independent; the components of one model are not, and they need operator-level, not job-level, decisions.
The gap. Nothing jointly decided per-operator parallelism and dependency-respecting schedule for a heterogeneous model on a shared cluster.

Where Spindle sits

The execution planner, run once per workload, and the runtime

Graph contraction new

Fuse chains of identical operators into MetaOps; assign each a MetaLevel by breadth-first depth so operators at one level are independent.
↓ →
Scalability estimator new

Profile a few GPU counts per MetaOp and fit a piecewise alpha-beta curve of time against GPUs; about five minutes.
↓ →
Resource allocator new

Per MetaLevel, solve the continuous malleable scheduling problem by bisection, then discretize each allocation into two valid integer GPU counts.
↓ →
Wavefront scheduler new

Greedily pack MetaOp slices into waves that use all devices with equal duration; concatenate levels.
↓ →
Device placement new

Map waves to GPUs preferring NVLink islands for heavy flows and balancing memory, with backtracking on out-of-memory.
↓ →
Runtime engine new

Instantiate MetaOps per device, insert send/receive between waves, build parameter device groups, and train wave by wave.

Everything is new. Planning takes under three seconds after profiling and is redone only when the task mix changes.

The idea

The optimization problem. Spindle treats all tasks as one directed acyclic computation graph. Each operator must be given an allocation of n GPUs and a start time s, minimizing the latest completion time, subject to never exceeding N GPUs at any instant and never starting an operator before its predecessors finish. Solving this directly over thousands of operators is hopeless, so Spindle simplifies in three moves.

Contract and level. Consecutive operators with a single edge between them and the same type and input size are the same workload, so they are fused into a

MetaOp

A run of consecutive identical operators (for example, twelve stacked Transformer layers on the same input shape) treated as one schedulable unit with one scaling curve.

. The 4-task Multitask-CLIP example shrinks to seven MetaOps. A breadth-first traversal assigns each MetaOp a

MetaLevel

The dependency depth of a MetaOp; MetaOps at the same level have no data flow between them and can be scheduled without ordering constraints.

, which removes the dependency constraint inside a level and lets each level be planned on its own, with levels concatenated afterwards.

Measure how each MetaOp scales. The standard alpha-beta communication model does not fit these workloads, because different kernels get invoked as the per-GPU slice changes. Spindle profiles a handful of GPU counts per MetaOp and fits a piecewise function whose coefficients change between ranges. Figure 4 shows why this matters: a vision MetaOp speeds up almost linearly with GPUs while a text MetaOp barely speeds up at all, so giving them equal shares is wrong in both directions.

Allocate by solving a relaxed problem exactly. Within one level, if GPUs and operators were continuously divisible, the problem is

malleable project scheduling

A scheduling problem in which each job's duration depends on how much of a divisible resource it receives, and the resource is shared among jobs running concurrently.

, and a classical theorem says the optimum has every MetaOp start at time zero, run all its operators, and finish at the same moment. That moment is found by bisection: guess a completion time, invert each scaling curve to find the GPUs each MetaOp needs to finish by then, and check whether the total is N. Then each real-valued allocation is written as two valid integer allocations. MetaOp 2 in the paper’s example, optimal at 1.5 GPUs over 12 operators, becomes 8.4 operators on 2 GPUs plus 3.6 on 1, rounded to 8 and 4, which preserves both the total work and the completion time. Valid means practical: a data-parallel MetaOp’s count must divide its batch, a tensor-parallel one must be a multiple of its degree.

Pack into waves. Discretization breaks the everyone-finishes-together property, so a

wave

The smallest scheduling unit: a set of MetaOp slices that run concurrently on disjoint, fixed groups of devices for one aligned time span, with data exchanged only at wave boundaries.

is introduced. Algorithm 1 builds waves greedily. Propose the allocation tuples that fill the most devices; if devices remain, extend the allocation of the MetaOp with the most remaining work; then align every proposed slice to the shortest one’s duration by scheduling only as many of its operators as fit, leaving the rest for later waves. In the paper’s six-wave example, wave 2 runs MetaOps 1, 2 and 4 on 4, 2 and 2 GPUs, but schedules only 1, 2 and 3 of their operators respectively so the wave ends together. Waves are at most twice the number of MetaOps, so scheduling is linear.

Place for locality and memory. A

device island

A set of GPUs joined by high-bandwidth links such as NVLink, typically the eight GPUs in one node.

is preferred for each MetaOp and for the data flows between waves, with the heaviest flows given first claim on intra-island placement. Because devices now hold different MetaOps, memory is uneven; placement estimates each MetaOp’s memory, prefers the emptiest device, co-locates MetaOps that share parameters, and backtracks earlier waves if a later one would run out of memory.

Run it. Each device instantiates its MetaOps, transmission operators (copy, shard, concat, send, receive) are inserted at wave boundaries for activations and gradients, and every shared parameter is assigned a device group over which its gradient is synchronized after the backward pass. Training proceeds forward wave by wave, then backward wave by wave, then group-wise synchronization.

Tracing the planner on the paper’s three-MetaOp example

Four GPUs, three MetaOps in one MetaLevel (Figure 5a)

MetaOp 1 has 8 operators, MetaOp 2 has 12, MetaOp 3 has 6. Allocations are written as (GPUs, operators). C* is the relaxed optimum. Numbers follow the paper's figure; scaling curves are illustrative.

Step 1 of 6

Fit scaling curves

Profiling gives each MetaOp a curve T(n). Suppose MetaOp 1 scales well, MetaOp 2 moderately, and MetaOp 3 is cheap and scales poorly. These curves, not raw sizes, decide the allocation.

T1(n), T2(n), T3(n) from ~5 minutes of profiling | piecewise alpha-beta fit

Why it works

Operator granularity exposes the real heterogeneity. Because allocation is decided per MetaOp from its own scaling curve, a light audio encoder gets four GPUs while a heavy vision encoder gets many, whereas task-level schemes give both the same share. The case study shows every device and every MetaOp near full utilization.
Levels turn one hard problem into several easy ones. Removing dependencies inside a level makes the relaxed problem a classical one with a closed-form optimality condition, so bisection finds the optimum instead of a heuristic guess.
Bi-point discretization preserves the optimum. Splitting each fractional allocation into two integer tuples with the same total time keeps the relaxed solution’s value; only rounding operator counts adds error, measured at under 7 percent.
Waves convert a schedule into something a runtime can execute. Fixed device groups and boundary-only communication mean no dynamic migration, and aligning spans removes the idle tails that discretization would otherwise create.
Placement keeps the extra communication small. Inter-wave send and receive is Spindle’s added cost. With island-aware placement it stays under 6 percent of step time; with naive sequential placement it reaches 27 percent.
Planning is cheap enough to redo. Under three seconds per plan, so when tasks join or leave, Spindle replans and keeps its advantage on dynamic workloads.

What it costs

Profiling and a model of the workload. Five minutes of profiling per model, and a fitted curve whose accuracy bounds the plan’s quality. Kernels that change behaviour at untested scales will be mispredicted.
Inter-wave communication. Activations and gradients cross wave boundaries between device groups; the paper keeps this to 6 percent but it is not zero, and it depends on the cluster’s island structure.
Rounding losses. Integer operator counts and resource extension in waves push the plan up to 7 percent above the relaxed optimum.
Parameter synchronization across groups. Shared parameters live on several device groups and are all-reduced after each step; 5 to 15 percent of step time in the breakdown.
Heuristic placement. Device placement is a constrained-depth recursive search with heuristics and backtracking, not an optimal solver; its time varies with the model.
Scope. Evaluated on models up to 9.25 billion parameters on 64 GPUs; 30-billion and 70-billion results on 256 GPUs are simulated, not measured. The baselines were tuned by hand because they cannot plan automatically for these models.

The evaluation

An 8-node cluster, each node with eight NVIDIA A800 80 GB GPUs on NVLink, nodes linked by 400 Gbps InfiniBand. Three models: Multitask-CLIP (1.20 billion parameters, 6 modalities, up to 10 contrastive tasks, ImageBind structure), OFASys (0.66 billion, 6 modalities, up to 7 tasks, encoder-decoder LM), and Qwen-VAL (9.25 billion, 3 modalities, 3 tasks, decoder-only LLM). Baselines Megatron-LM and DeepSpeed, hand-tuned per workload, plus Spindle-Optimus and DistMM-MT. Average of 100 iterations.

Iteration-time speedup over DeepSpeed, and where it comes from

Property	Megatron-LM / DeepSpeed	Spindle-Optimus (task-level)	DistMM-MT (intra-task)
Best speedup, all workloads	1.0
Multitask-CLIP, 8 / 16 / 32 GPUs	1.0	—	—
Scaling 16 to 32 GPUs, 4-task Multitask-CLIP		—	—
Qwen-VAL 9.25B, 32 / 64 GPUs	1.0	—	—
Gap to theoretical optimum	—	—	—
Inter-wave send/receive overhead	none	—	—
Planning time	manual tuning	—	—

From Figures 8 to 12 of the paper. Appendix results: on dynamic task sets Spindle keeps the shortest total training time; simulated on 256 GPUs, Spindle is 1.34x and 1.36x over DeepSpeed for 30B and 70B Qwen-VAL versus at most 1.09x for the other baselines; on a single-task workload Spindle still beats SOTA by up to 48 percent; per-device memory is lower and more balanced than the baselines because only devices that activate an operator store its parameters.Underlined values explain themselves when clicked.

In an actual implementation

Reading the case study

Figure 9 looks at 4-task Multitask-CLIP on 16 GPUs over one iteration. DeepSpeed’s cluster utilization oscillates as heavy and light components alternate. Spindle-Optimus starts high but decays as light tasks finish and their GPUs go idle. DistMM-MT is better within tasks but cannot balance across them. Spindle stays near the top for the whole iteration, and the spider charts show every one of the 16 devices and every MetaOp close to the outer edge. The time breakdown in Figure 10 explains the totals: forward and backward are 80 to 95 percent of a step in every system, parameter synchronization 5 to 15 percent, and Spindle’s added send and receive 1.7 to 8.7 percent depending on workload, versus 10 to 27 percent if placement ignores islands.

Where this stops working

Spindle’s plan is only as good as its scaling curves, and those come from a few profiled points per MetaOp; a kernel whose behaviour changes at an unprofiled scale will be mis-sized. The per-level decomposition ignores any slack between levels, and the wave abstraction forbids overlapping a MetaOp’s computation with the next wave’s communication, so some pipelining is left on the table. The largest measured configuration is 9.25 billion parameters on 64 GPUs; the 30B and 70B numbers on 256 GPUs are from a simulator, and the authors note that sub-10B multi-modal models are the common deployed case. Baselines were tuned by hand, so their numbers reflect the authors’ tuning effort. Finally, the whole design targets models with several dependent, unequal components; a single homogeneous LLM gains nothing from it.

Background

Why the relaxed problem has such a clean answer

Malleable project scheduling with continuously divisible resources was studied by Weglarz in the early 1980s. If every job’s duration is a positive, non-increasing function of the resource it gets, and jobs are independent, then the best way to share the resource is to give each job a constant amount so that all of them finish at the same instant: any job finishing early has resource that could have shortened the longest one. That turns a scheduling problem into a one-dimensional search for the common finish time, which is exactly what Spindle’s bisection does. The engineering is in making reality fit the theorem: contracting operators so there are few enough jobs, levelling so they are independent, fitting curves that are non-increasing, and discretizing without breaking the equal-finish property.

If you remember only three things

1. Multi-task multi-modal models are built from components that differ by an order of magnitude in cost and in scalability, share layers, and depend on one another. Training systems built for one homogeneous model run them one at a time across the whole cluster, so light components leave GPUs idle; task-level allocation fixes only half the heterogeneity and idles GPUs as short tasks finish; single-task multi-modal systems never optimize across tasks.
2. Spindle contracts identical operator chains into MetaOps, groups them into dependency levels, fits per-MetaOp scaling curves from five minutes of profiling, and for each level solves the relaxed malleable scheduling problem exactly by bisection (all MetaOps finish together) before discretizing each allocation into two valid integer GPU counts. A greedy wavefront scheduler packs MetaOp slices into device-filling, duration-aligned waves; placement keeps heavy flows inside NVLink islands and balances memory; the runtime inserts inter-wave transfers and synchronizes shared parameters within device groups.
3. On an A800 cluster, Spindle trains Multitask-CLIP, OFASys and 9.25B Qwen-VAL up to 71 percent faster than Megatron-LM and DeepSpeed, up to 44 and 59 percent faster than task-level and intra-task baselines, within 7 percent of the theoretical optimum, with inter-wave communication under 6 percent and plans in under three seconds. Costs are profiling, rounding loss, cross-group parameter synchronization and heuristic placement; the largest measured run is 64 GPUs, with 256-GPU results simulated.

Check yourself

Check your understanding

1. Why does giving an entire cluster to one task at a time, as Megatron-LM and DeepSpeed do, waste GPUs on multi-task multi-modal models?
2. What property does the relaxed allocation problem have that makes it solvable by bisection?
3. Predict what happens if Spindle skipped the wavefront scheduler and simply started every discretized allocation tuple at time zero.
4. Why is device placement a separate step, and what does ignoring it cost?
5. Which result most directly supports the claim that Spindle's advantage grows with model complexity?

Sources

The paper is “Spindle: Efficient Distributed Training of Multi-Task Large Models via Wavefront Scheduling” by Yujie Wang, Shenhan Zhu, Fangcheng Fu, Xupeng Miao, Jie Zhang, Juan Zhu, Fan Hong, Yong Li and Bin Cui, Peking University, Purdue University and Alibaba Group, published at ASPLOS 2025, Volume 2, DOI 10.1145/3676641.3715992. This explainer was written from the full text of the arXiv preprint (2409.03365, version 3, February 2025), including its appendix. Numbers are taken from Figures 1, 4, 5 and 8 to 16 and Tables 1 and 2.