Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training

Making a giant model train fast means rewriting it — fusing operators, sharding weights across GPUs, checkpointing activations — and every such change is stitched into the model's own source, so the code that a researcher wrote and the code that runs efficiently become the same tangled thing, unusable on the next model or the next machine; Slapo (Cornell and Amazon) borrows a decades-old compiler idea and writes the optimizations as a separate schedule that names parts of the untouched model and transforms them.

Why this problem is worth solving

The optimizations that make training fast are surgery on the model itself.

There is a well-worn idea for exactly this shape of problem.

What people did before, and where it broke

• Hand-optimise the model in a distributed-training framework. The dominant practice: frameworks provide manually parallelised implementations of common architectures, tuned hard. Where it broke. The optimizations are intrusive edits, so a new model — even one barely different — cannot reuse them, the result runs only on the platform it was tuned for, and separate copies must be maintained for single-device and multi-device use.
• Trace the whole model into a static graph and let a compiler optimise it. The automatic route: capture the entire computation as one graph and apply fusion, simplification and even parallelism. Where it broke. Everything must be statically analysable, which forbids the dynamic constructs — in-place updates, certain control flow, the collective operators central to distributed training — that real training code uses, so the model must often be rewritten to be traceable at all.
• Accept the flattened graph the compiler produces. A consequence of whole-graph tracing: the layer hierarchy a developer built collapses into a single flat dataflow graph, with operators decomposed into smaller pieces. Where it broke. The optimised model no longer resembles the original, so when performance or convergence goes wrong there is nothing recognisable to debug against.
• Expose tuning knobs by editing the model to surface them. To find the best configuration — how many pipeline stages, which activations to checkpoint — developers modify the model to make those choices adjustable. Where it broke. This ties the definition even more tightly to the optimizations and is tedious and error-prone, precisely because the two are entangled.
• The gap. One approach keeps the model readable and pays in optimizations that do not transfer; the other automates the optimizations and pays in a model that is neither traceable without rewriting nor debuggable afterward. Nobody had the separation that tensor compilers had already shown works one level down.

How a model gets faster without being rewritten

The idea

Lifting a proven idea one level up is the whole move, and the level is what makes it work. Tensor compilers separated an operator’s arithmetic from its schedule and got portability and auto-tuning for free; the observation here is that a model’s definition and its training optimizations stand in exactly the same relation, one level up. The definition says what the model computes; the schedule says how to run it fast on this hardware — which modules to shard, which patterns to fuse, where to split the pipeline. Written apart, the model stays the clean thing a researcher wrote and reruns unchanged on a single device or a different platform, while the schedule carries everything platform-specific and can be swapped, tuned, or discarded. The payoff shows up as effort: the schedules that deliver the speedups are on the order of tens of lines against model definitions of over a thousand, because the schedule references parts of the model rather than reproducing them. Taking a separation that solved a tension at the operator level and applying it to the model level is a direct sibling of the tensor-program work covered elsewhere on this site, which pushes in the opposite direction — arguing that the declarative loop schedule is too blunt and giving the programmer finer control — where this paper argues the schedule idea was not being applied high enough.

Tracing only where you must is what lets the schedule coexist with everything dynamic. The reason whole-graph compilers are so disruptive is that they demand the entire model be static, which collides with the in-place updates, data-dependent control flow and collective communication operators that distributed training is built from — and then, having traced everything, they flatten the hierarchy so nothing is recognisable afterward. The design here refuses the all-or-nothing choice by splitting the primitives in two. Module-level transformations — replacing a submodule with a faster kernel, sharding a weight, checkpointing — do not alter the forward computation, so they need no static graph at all and can be applied directly to the dynamic model. Only the transformations that genuinely rewrite computation, operator fusion and pipeline splitting, require a graph, and those trace just the one module they act on. So the traced region is as small as the optimization demands, the rest of the model remains dynamic and debuggable, and the layer hierarchy is preserved because scheduling addresses modules by their path in the tree. Meeting a general workload where it is by transforming only the part that needs it, rather than forcing everything into one rigid form, is the same accommodation the dynamic-shape compiler work covered elsewhere on this site makes when it partially lowers portions of a program and leaves the rest for later.

Separating the schedule turns tuning and portability from manual work into structure. Once the optimizations live apart from the model, several things that were laborious become almost automatic. The configuration choices inside a schedule — how many pipeline stages, what fraction of activations to checkpoint — are just variables, so declaring them as a search space and handing that to an auto-tuner replaces the manual edit-and-benchmark loop, which matters because the evaluation shows no single parallelism strategy is best across GPU counts and models. Portability comes for free in the same way: the scheduled model runs on plain PyTorch by default, and thin adapters let the same schedule target the specialised runtimes of existing distributed frameworks by handling their input and output conventions automatically. And extensibility is built in — a user can register a new optimization as a primitive that cooperates with the tuner and verifier, which the paper reports taking on the order of an hour and tens of lines. That once the optimizations are named and separated a search procedure can explore them mechanically is the same benefit the Bayesian autotuning work covered elsewhere on this site pursues when it treats a compiler’s interacting knobs as a space to be searched rather than a configuration to be hand-set — here the separation of schedule from definition is what turns the knobs into a space at all.

Following one attention layer from clean to optimised

This is a module replacement addressed by path, so it applies to the dynamic model directly. The reverse operation later — splitting the fused output back into three — is a computation rewrite that does need tracing.

Why it works

• The definition-and-schedule relation is the same one tensor compilers already solved. A model’s arithmetic is what a researcher wants to keep readable, while its parallelism and kernel choices are platform-specific performance work → and that is exactly the split tensor compilers made between an operator’s definition and its schedule → so lifting the same separation up a level gives the model the portability and tunability the operator already had.
• Most optimizations do not change what is computed, only how modules are arranged. Replacing a submodule, sharding a weight and checkpointing all leave the forward computation intact → so they can be applied to the dynamic model directly, with no tracing and none of its constraints → which is why the disruptive static-graph requirement is confined to the few primitives that truly rewrite computation.
• Tracing by need keeps the traced region minimal. Only operator fusion and pipeline splitting must manipulate the forward pass, and only within specific modules → so the schedule traces just those modules and leaves the rest dynamic → preserving debuggability everywhere the optimization does not reach and avoiding the whole-model flattening.
• Naming modules by path preserves the hierarchy the developer relies on. Whole-graph compilers flatten the tree because they optimise across it globally → whereas addressing a module by its place in the tree means the tree is still there after scheduling → so a developer can locate a layer, apply a change, and debug against a structure they recognise.
• A separate schedule is a search space by construction. The configuration choices are entangled in the model under manual optimization, so tuning means editing → but in a separate schedule they are named variables → so an auto-tuner can explore them mechanically, which is necessary because the best parallelism strategy varies with model and GPU count.

What it costs (trade-offs and scope)

• It orchestrates existing optimizations rather than inventing them. Every speedup comes from a kernel, a parallelism scheme or a checkpointing strategy that already existed; the contribution is making them composable and transferable, not faster in themselves. The paper is clear that it schedules hand-crafted optimizations systematically.
• The headline speedup is against out-of-the-box baselines. The up-to-2.92-times figure compares against the default configurations of the baseline frameworks, and much of it comes from the baselines not supporting a model well or at all. Where a baseline is strong and runs the same underlying scheme, the margin is far smaller.
• On the fair comparison the gains are modest. Against one baseline restricted to the three models it officially supports, the tensor-parallel speedup is 1.02 to 1.46 times, and the paper notes it does not significantly outperform that baseline on two of the three, attributing the difference partly to the model implementations rather than the scheduling.
• The graph tracer is inherited and limited. Tracing uses PyTorch’s own mechanism, which cannot capture certain control flow and data structures and flattens what it does capture — so any computation-rewriting primitive is bounded by what the tracer supports, and the paper had to build on top of it rather than around it.
• The schedule can be wrong, and the verifier is rule-based. Because primitives are flexible, a developer can schedule incorrectly — a shard without its aggregation, a layout mismatch — and the safety net is a set of predefined rules per primitive plus a numerical check, which catches the anticipated mistakes rather than guaranteeing correctness in general.
• Portability to a specialised runtime needs a hand-written adapter. Running a schedule on a distributed framework’s own pipeline runtime requires a per-framework dialect handling that runtime’s input and output conventions; two are provided, and a third target would need another.
• Auto-scheduling is future work. The tuner explores a space the user declares; it does not discover the optimizations or the space itself, which the paper lists as a planned extension modelled on prior auto-schedulers.
• The scope is a familiar model family on one accelerator vendor. Transformers and a convolutional network, all from standard hubs, trained on one cloud GPU generation. The claim of platform portability is architectural rather than demonstrated across vendors.
• Extensibility is measured on the authors’ own additions. The reported hour-and-tens-of-lines cost of a new primitive comes from the authors adding primitives they designed, and explicitly excludes comprehensive testing across models.
• The abstraction is PyTorch-shaped. The whole approach leans on PyTorch’s dynamic graph and module hierarchy; it is a schedule language for that framework’s style, not a framework-neutral result.
• Correctness rests on the transformations being semantics-preserving, which for sharding and fusion depends on the developer supplying any needed layout logic — the verifier checks structure and numerics but the burden of a correct transformation is shared.

The evaluation

Check yourself

Sources

The paper is “Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training” by Hongzheng Chen, Cody Hao Yu, Shuai Zheng, Zhen Zhang, Zhiru Zhang and Yida Wang, published at ASPLOS 2024, pages 1095-1111, DOI 10.1145/3620665.3640399. The authors are at Cornell University and Amazon.

This explainer is written from the full arXiv text, arXiv:2302.08005. Everything above comes from the paper’s own sections: the four challenges motivating the work — generality and ease of tuning for manual optimization, programmability and debuggability for compiler-based optimization; the background on high-performance kernels, activation checkpointing and the three forms of parallelism, and the worked BERT attention example showing how fusing the query-key-value projections, swapping in an efficient attention kernel, sharding for tensor parallelism and partitioning for pipeline parallelism are each intrusive edits; the design of Slapo, covering its hierarchical module representation, the split between module-level primitives that need no static graph and computation-level primitives that do, the trace-by-need approach, the parameter sharding and synchronisation primitives, the operator fusion and pipeline partitioning primitives, the extensible custom-primitive mechanism, the auto-tuner and its search-space construction, the framework dialects for two distributed runtimes, and the rule-based plus numerical verifier; the implementation notes on building atop PyTorch and its tracer’s limitations; and the evaluation covering the single-machine and multi-machine throughput comparisons against the two baselines, the tensor-parallel margins on the officially supported models, the line-of-code comparison, the ablation over progressively applied primitives, the auto-tuning study, and the development cost of adding new primitives.