PartIR: Composing SPMD Partitioning Strategies for Machine Learning

Sami Alabed, Daniel Belov, Bart Chrzaszcz, Juliana Franco, Dominik Grewe, Dougal Maclaurin and 11 others

Training a giant neural network means splitting it across thousands of accelerator chips using a mix of strategies — sharding the data, the model's weights, and the optimizer state — and today's tools make engineers scatter cryptic sharding hints throughout their model code, which ties the code to one cluster and cannot be checked strategy by strategy; this work from Google DeepMind builds PartIR, which pulls the partitioning out of the model into a separate, ordered recipe of composable tactics that a compiler applies by principled rules, so each strategy can be verified on its own without a costly test run, and matches the performance of the state-of-the-art tool while being simpler and predictable.

The argument in one glance

Problem

Splitting a huge model across chips needs many strategies

A modern large neural network no longer fits on one accelerator, so training splits it across thousands of them using a combination of strategies — sharding the input data, splitting the model weights, and partitioning the optimizer state to save memory. Making this run correctly requires inserting exactly the right cross-device communication.
What people did

Annotate the model code with sharding hints

Partitioning tools let engineers avoid hand-writing low-level communication by adding sharding annotations to their model code. The tool reads the hints and figures out the communication automatically.
Why it fell short

Annotations tie code to one cluster and cannot be checked

Burying sharding hints in the model has three costs. The annotated code is tied to one cluster's shape, so a model pretrained on one system and fine-tuned or deployed on another must be re-annotated. Each strategy cannot be verified on its own, forcing slow trial-and-error profiling. And resolving conflicts between strategies means placing annotations just so, by guesswork.
This paper's idea

Make partitioning a separate, composable schedule (PartIR)

PartIR pulls the partitioning out of the model into a separate schedule — an ordered sequence of tactics, each a parallelism strategy, composed in any order and never undoing earlier ones. A compiler applies each tactic by principled rewrite rules derived from the operations' algebra, rewriting the program after each so its communication can be verified analytically, without a test run. It matches the state-of-the-art tool's performance while being simpler and predictable.

Why this problem is worth solving

Scaling a model past one chip is now routine, and doing it well is a tangle of interacting choices.

Modern neural networks have grown far faster than the memory and speed of any single accelerator, so training runs across many chips at once. The dominant way to do this is called SPMD, for Single Program, Multiple Data: every device runs the same program but on a different slice of the data. Deciding how to slice the work is called partitioning, and placing a piece of a particular tensor on a particular device is called sharding. Several strategies are combined in practice. Data parallelism shards the input batch across devices while replicating the model, so each device trains on different examples. Model parallelism shards the model’s weight matrices themselves across devices, which is needed when the model is too big to replicate. Optimizer sharding splits the optimizer state, and sometimes the gradients and parameters too, across devices to cut peak memory. Because the devices hold different pieces, they must exchange data at the right moments using collective communication operations — for example an all-reduce that sums a value across all devices, or an all-gather that reassembles a full tensor from its shards. The devices are arranged logically as a mesh with named axes reflecting the network, and a large model is trained by combining several of these strategies along different mesh axes at once. Getting the combination right, and inserting exactly the needed collectives, is intricate and central to training efficiency.

The tools that automate this make a choice that quietly causes pain.

Writing the low-level communication operations by hand is error-prone and non-portable, so partitioning tools do it for the engineer. The prevailing approach, used by widely-deployed systems, is to have the engineer annotate the model code with sharding hints — small markers on tensors saying how they should be split — and the tool then propagates those hints and inserts the collectives. This works, but embedding the sharding decisions inside the model code has three real costs. First, portability: the annotations bake in the shape of one particular cluster, so a model pretrained on one system that then needs fine-tuning on a second and deployment on a third has to be re-annotated for each, and the model code no longer outlives the hardware it was written for. Second, verifiability: because the strategy is tangled into the model, an engineer cannot check each parallelism strategy on its own; to find out what the partitioner actually did, they must run the whole thing and profile it, a slow trial-and-error loop. Third, conflicts: when two strategies want to shard the same tensor dimension in incompatible ways, the tool needs the annotations placed just so to resolve the conflict, which becomes a guessing game. These frictions grow with the model and the number of strategies, which is exactly where scaling is hardest.

What people did before, and where it broke

Hand-write the collective communication. Manually insert the all-reduces, all-gathers, and so on for each parallel strategy. Where it broke. It is error-prone and non-portable; every change of model or cluster means rewriting delicate low-level distributed code.
Annotation-based partitioners (GSPMD and similar). Let the engineer mark tensors in the model with sharding hints and have the tool propagate them and insert collectives. Where it broke. The annotations tie the model to one cluster, cannot be verified strategy by strategy without profiling, and need trial-and-error placement to resolve conflicts between strategies.
Fully automatic partitioners. Search for a good partitioning automatically. Where it broke. They add flexibility but are unpredictable and slow, and users still want to combine them with known, trusted strategies like data parallelism for speed and guarantees.
The gap. No system let engineers compose parallelism strategies, manual or automatic, separately from the model code, with predictable rules and the ability to verify each strategy independently. Providing that, by making partitioning a composable schedule over the program, is what PartIR does.

How PartIR works

From sharding hints in model code to a separate, verifiable schedule

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Separate the model from how it is partitioned. PartIR’s founding move is to split two things that annotation-based tools fuse: what the model computes, and how that computation is spread across devices. The model code stays clean and hardware-agnostic; the partitioning lives in a separate artifact, so the same model can be re-partitioned for a different cluster by swapping that artifact rather than editing the model. Pulling the description of a computation apart from the details of how it maps onto particular hardware, so the first can outlive and be retargeted independently of the second, is the same layering the fixed-point compiler work covered elsewhere on this site relies on, lifting a program into one representation and then lowering it separately for each target rather than entangling the two.

Write the partitioning as an ordered, verifiable schedule. PartIR expresses the partitioning as a schedule: a sequence of tactics, each one a parallelism strategy, applied in order, where no tactic can undo the sharding a previous one chose. After each tactic PartIR rewrites the program and makes the communication explicit, so the engineer can read off the collectives and estimate the performance of that strategy analytically — confirming, say, that data parallelism produced exactly one all-reduce per parameter — without ever running the model. Composing an optimization as an ordered recipe of steps whose effect can be estimated by analysis rather than measured by an expensive run is the same discipline the schedule-drafting work covered elsewhere on this site follows, proposing a schedule and checking its predicted performance with a model instead of profiling every candidate.

Propagate by principled rules, not by guessing. Rather than resolve sharding by cost-based heuristics or by asking the user to place annotations until conflicts disappear, PartIR propagates each sharding decision across the whole program using rewrite rules derived from the algebraic meaning of each operation, and it resolves conflicts by the order of the tactics. The result is deterministic and predictable, and the critical transformation is even formally verified. Replacing an error-prone pile of ad-hoc, trial-and-error decisions with a small set of principled rules that behave the same way every time is the same shift the copy-on-write correctness work covered elsewhere on this site makes, trading a fog of special-case patches for one well-defined rule that a system can apply predictably and safely.

Following a model being partitioned

Partitioning a two-matrix-multiply model with a schedule

Tracing one schedule shows how tactics compose and are verified step by step.

Step 1 of 5

Write the model, with no sharding hints

An engineer writes a small model: multiply the input by one weight matrix, then multiply the result by another. There are no sharding annotations anywhere in the model code.

The model is just the computation, kept independent of any cluster.

Why it works

Combining strategies is unavoidable but complex. Because a large model exceeds one device and needs data, model, and optimizer sharding at once → the partitioning becomes a tangle of interacting decisions → so a way to compose strategies cleanly is essential to scaling.
Annotations bury the strategy in the model. Because existing partitioners require sharding hints inside the model code → the model is tied to one cluster and no strategy can be checked alone → so decoupling the strategy from the model restores portability and testability.
A schedule makes each strategy a separate, verifiable step. Because the partitioning is an ordered sequence of tactics that rewrite the program one at a time → the collectives after each tactic can be estimated analytically → so the engineer verifies performance without profiling and reorders tactics to resolve conflicts predictably.
Rules from semantics beat cost heuristics. Because sharding is propagated by rewrite rules derived from each operation’s algebra rather than cost-based guessing → the result is deterministic and the same every time → so there is no trial-and-error, and no heuristic quietly degrades performance.
Decoupling and verification cost nothing in speed. Because the partitioned program uses the same collectives an expert would write → its device utilization and memory match the state of the art → so PartIR matches the leading tool within about one percent while being simpler and predictable.

What it costs (trade-offs and scope)

The engineer must write a schedule. PartIR removes annotations from the model but requires a separate partitioning schedule. This is a cleaner artifact, but it is still work, and it asks the engineer to think about tactics and their order.
Ordering matters and must be understood. Because tactics never undo earlier decisions, the order in the schedule determines how conflicts resolve. This is predictable, but it means the engineer has to reason about tactic order rather than have a tool guess for them.
It is a compiler system to build and maintain. PartIR is a stack of compiler passes and intermediate representations built on a common compiler infrastructure. Getting its benefits requires that infrastructure, not a small library drop-in, though it works across multiple frontends and backends.
Automatic tactics inherit search costs. PartIR can compose automatic partitioning tactics as well as manual ones, but where it uses automatic search, it takes on that search’s cost and any unpredictability, mitigated by combining it with fast, known manual tactics.
Matching, not beating, the state of the art on raw performance. PartIR’s performance is on par with the leading annotation-based tool, within about one percent; its advantage is in simplicity, portability, predictability, and verifiability, not in a large speedup.
The predictability rests on the rule set. The guarantee that partitioning is deterministic and analyzable depends on the algebraic rewrite rules covering the operations in use; an operation without well-defined partitioning rules would fall outside the clean story.

The evaluation

Three ways to partition a model across devices

Property	Hand-written collectives	Annotation partitioners
Model code carries sharding	Not applicable
Portable across clusters		No, re-annotate
Verify each strategy without running	No
Resolving strategy conflicts	Manual
Performance (utilization and memory)	Expert-dependent	State of the art

Three ways to partition a neural network across accelerators, evaluated on Transformer models trained on tens to over a hundred TPUs and on GPUs. Hand-writing the collective communication is error-prone and non-portable. Annotation-based partitioners like GSPMD embed sharding hints in the model, which ties the code to one cluster, cannot be verified strategy by strategy without profiling, and needs trial-and-error to resolve conflicts. PartIR instead separates the model from a schedule of composable tactics, rewriting the program after each tactic so its collectives can be verified analytically, and propagating sharding by principled rewrite rules whose critical transformation is formally verified. Across the tested models and hardware, PartIR matches GSPMD's model-FLOPS utilization and memory usage to within about one percent, while offering a simpler API, portability, and per-step verification that annotation tools do not; the authors also found that PartIR's rule-based design avoids a class of silent performance-degrading bugs. PartIR's learnings were folded into the open-source Shardy project.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

PartIR is a clean, practical compiler contribution from a team that trains models at real scale. Its core insight borrows a lesson from image-processing compilers — separate the algorithm from the schedule — and applies it, for the first time, to partitioning whole neural-network programs across devices. The payoff is threefold: the model code stays free of sharding annotations, so it is portable and outlives any cluster; the partitioning is an ordered, composable schedule of tactics whose effect can be verified analytically after each step, so engineers stop debugging by profiling; and sharding propagates by principled rewrite rules rather than cost heuristics or trial-and-error, so the result is predictable and the critical transformation is formally verified. Importantly, this is not a raw-speed win: PartIR matches the state-of-the-art annotation-based tool within about one percent on utilization and memory. Its advantages are simplicity, portability, predictability, and the ability to verify each strategy — including catching a class of silent performance bugs that heuristic conflict-resolution can introduce. Read it with its scope. The engineer still writes a schedule and must reason about tactic order; PartIR is a full compiler stack, not a small library; automatic tactics inherit their search costs; and the predictability rests on the algebraic rules covering the operations used. PartIR is now part of the open-source Shardy project. The durable idea is that when a hard task fuses two concerns — here, what to compute and how to distribute it — pulling them apart into a program and a composable, verifiable schedule makes the whole thing tractable.

Background

Why separating the what from the how keeps paying off

It helps to see PartIR as a fresh instance of one of the most durable ideas in compiler design: separate what a program computes from how it is executed, and make the how a first-class, composable, checkable thing. In image processing, this separation of algorithm from schedule transformed how high-performance kernels are written; PartIR carries the same idea into distributed model training, where the how is the partitioning across a cluster. The synthesis-based hardware-mapping work covered elsewhere on this site pursues a kindred goal in another domain, replacing hand-written, per-target mapping rules with a principled system that decouples the intent from the target and can be checked for correctness, so that a task once done by fragile manual effort becomes systematic. The common thread is that entangling the essential computation with the messy details of how it runs on particular hardware makes both harder to get right, and that pulling them apart — with the execution side expressed as something you can compose, retarget, and verify — unlocks portability and confidence. As models and hardware both grow more varied, the value of a clean seam between the model and its distribution only rises, because it lets the model outlive the machine and lets the distribution be reasoned about on its own.

If you remember only three things

1. Training a large neural network requires splitting it across many accelerators with a combination of parallelization strategies — data parallelism (shard the batch), model parallelism (shard the weights), and optimizer sharding (shard the optimizer state to save memory) — in the single-program-multiple-data style, which means inserting exactly the right cross-device communication collectives. Existing tools automate this but make engineers annotate the model code with sharding hints, which ties the code to one cluster (so it must be re-annotated to move), prevents verifying each strategy on its own (forcing slow profiling), and needs trial-and-error to resolve conflicts between strategies.
2. PartIR separates the model from its partitioning. The model code carries no sharding annotations; the partitioning is a separate schedule — an ordered sequence of tactics, each a parallelism strategy, composed in any order and never undoing earlier sharding decisions. A compiler applies each tactic by rewrite rules derived from the operations' algebraic semantics (not cost heuristics), rewriting the program's intermediate representation after each tactic and making the communication explicit, so an engineer or simulator can verify that strategy's collectives and estimate its performance analytically without a test run. Conflicts are resolved predictably by tactic order, and the critical transformation is formally verified. It is an MLIR-based compiler usable from any frontend to any backend.
3. On Transformer models trained across tens to over a hundred TPUs and on GPUs, PartIR matches the state-of-the-art annotation-based tool (GSPMD) on model-FLOPS utilization and memory to within about one percent, while adding a simpler API, portability, per-step verifiability, and freedom from a class of silent performance-degrading bugs that heuristic conflict resolution can cause; its learnings were folded into the open-source Shardy project. Its scope: the engineer still writes a schedule and must reason about tactic order, it is a full compiler stack rather than a drop-in library, automatic tactics inherit search costs, and the performance is on par with rather than beyond the state of the art.

Check yourself

Check your understanding

1. Why does training a large neural network need several parallelization strategies combined?
2. What is the core problem with annotation-based partitioners that PartIR addresses?
3. How does PartIR let an engineer verify a parallelism strategy without running the model?
4. PartIR matches the state-of-the-art tool's performance to within about one percent. Why is it still valuable?

Sources

The paper is “PartIR: Composing SPMD Partitioning Strategies for Machine Learning” by Sami Alabed, Daniel Belov, Bart Chrzaszcz, Juliana Franco, Dominik Grewe, Dougal Maclaurin, James Molloy, Tom Natan, Tamara Norman, Xiaoyue Pan, Adam Paszke, Norman A. Rink, Michael Schaarschmidt, Timur Sitdikov, Agnieszka Swietlik, Dimitrios Vytiniotis, and Joel Wee, published at ASPLOS 2025, DOI 10.1145/3669940.3707284. The authors are at Google DeepMind.

This explainer is written from the full text of the paper, obtained from the authors’ arXiv preprint (2401.11202), which matches the ASPLOS 2025 proceedings. Everything above comes from the paper’s own sections: the introduction and background on SPMD partitioning for machine learning (data, model, and optimizer/fully-sharded parallelism, the collective communication operations all-reduce, all-gather, reduce-scatter, and all-to-all, and the device mesh abstraction), the shortcomings of user-driven annotation-based partitioners such as GSPMD (non-portability across topologies, the inability to verify strategies independently and the resulting profiling, and trial-and-error conflict resolution) and of fully automatic partitioners (flexibility but unpredictability and slowness); PartIR’s design as a decoupled, expressive, and predictable partitioner using a schedule of composable tactics that incrementally rewrite the program’s intermediate representation and never undo earlier decisions, its propagation by rewrite rules derived from algebraic semantics rather than cost heuristics with conflicts resolved by tactic order, and its MLIR-based compiler with the PartIR:Core dialect (functional tiling and reduction loops over an array intermediate representation) and PartIR:HLO dialect (explicit, optimized SPMD collectives), with the critical Core-to-HLO transformation formally verified; the worked example of partitioning a two-matrix-multiply program with batch, model, and fully-sharded data parallelism; and the evaluation on U-Net, a graph network simulator, and the T32 and T48 Transformer models across TPU and GPU configurations, showing PartIR matches GSPMD’s model-FLOPS utilization and high-bandwidth-memory usage within about one percent, that the number of collectives after each tactic matches analytical expectations, and that PartIR’s rule-based design avoids a class of silent performance-degrading bugs, with its learnings incorporated into the open-source Shardy project.