PrimePar: Efficient Spatial-temporal Tensor Partitioning for Large Transformer Model Training

Training a giant transformer across GPUs splits each operation's tensors across devices, but every system does the split only across space — different pieces on different devices — which forces costly all-reduce communication whenever pieces hold partial sums, and wastes memory replicating shared tensors; this work from the Chinese Academy of Sciences adds a time dimension, borrowing Cannon's systolic matrix-multiply so a device accumulates its partial sums locally across successive steps instead of reducing them over the network, cutting collective communication to as little as a fifth and lifting training throughput up to 1.68 times at 69 percent of the memory.

The argument in one glance

Problem

Splitting tensors across devices forces all-reduce and wastes memory

Transformer models have grown to hundreds of billions of parameters, so training them means spreading each operator across many GPUs by tensor partitioning — the umbrella over data and model parallelism. But every existing system distributes the split pieces only across space, one piece per device. Whenever the dimension being split is one that is mathematically summed over, each device ends up with a partial sum, and turning those into the real answer needs all-reduce, a collective that is expensive and hard to hide behind computation. Worse, any tensor left unpartitioned is copied onto every device, wasting memory that only grows as you add GPUs.
What people did

Search the spatial partition space for the best plan

The mature response is to treat partitioning as a search problem. Systems like Megatron-LM hand-craft good spatial splits for transformer operators, slicing the rows and columns of linear weights and the head dimension of attention. Alpa goes further and automatically searches a recursive space of spatial partitions for the plan with the least communication. Both are strong and reach similar performance, because they explore the same well-understood space of ways to lay operators out across devices.
Why it fell short

That whole space is spatial, so the two costs are baked in

The trouble is the space itself. As long as every sub-operator must live on a distinct device, partitioning a summed-over dimension unavoidably scatters partial sums that only all-reduce can gather, and unpartitioned tensors are unavoidably replicated. Measurements on Megatron-LM show all-reduce taking 38 to 65 percent of training time, and its memory drifting ever further above the ideal as parallelism grows. No amount of searching a purely spatial space escapes costs that are built into the space.
This paper's idea

Add a time dimension so partial sums accumulate locally

PrimePar widens the space with a second axis: distribute some sub-operators across time on the same device rather than across devices. Borrowing Cannon's systolic matrix-multiply, a device runs a short sequence of steps in which tiles shift in a ring and partial sums pile up locally — so the summed-over dimension is reduced in place, over time, with no all-reduce at all, and shared tensors are reused across steps instead of replicated. The only device-to-device traffic left is cheap ring communication that overlaps with computation.

Why this problem is worth solving

The technique everything here turns on is the general way to run one operator on many machines.

Tensor partitioning is the general way to run a single neural-network operator across many devices: you take the big multi-dimensional arrays it works on and slice them along one or more of their dimensions, giving each device a piece to compute. It unifies what people usually call data parallelism (slice the batch) and model parallelism (slice the weights) under one framework, and a good partition plan can be searched for automatically. The subtlety is which dimension you slice. Some dimensions of a matrix multiply are summed over — the shared inner dimension that gets contracted away — and if you split one of those, each device computes only a partial sum of the true result, so the pieces must be added back together across the network. Other tensors may not contain the dimension you split at all, in which case each device needs a full copy, and the same data is replicated everywhere. Both effects are the hidden price of deciding how to cut the tensors.

The specific operation that price is paid in is the one this paper works hardest to avoid.

All-reduce is a collective communication operation in which every device contributes a value and every device receives the sum of all of them. It is the natural way to finish a computation that has been split so each device holds only a partial sum — a running total that is correct only once everyone’s contributions are added. The catch is that all-reduce is expensive and sits squarely on the critical path: the next step cannot proceed until the full sum is known everywhere, so unlike simple point-to-point messages it is genuinely hard to overlap with useful computation. In large transformer training it can consume well over half of the wall-clock time, and its cost rises with the number of devices taking part. Avoiding it, rather than merely making it faster, is therefore one of the biggest levers on training throughput.

What people did before, and where it broke

Hand-crafted spatial partitions (Megatron-LM). Expert-designed splits for transformer operators — row and column of the linear weights, head dimension of attention — mapped across devices. Where it broke. They are carefully tuned but still purely spatial, so any summed-over split still ends in all-reduce and any unpartitioned tensor is still replicated; the costs are structural, not a tuning mistake.
Automatic search over the spatial space (Alpa). Formalize a recursive space of spatial partitions and search it, as an integer linear program, for the plan with the least communication. Where it broke. It searches the very same spatial-only space, so it lands at the same wall as Megatron-LM, and the integer program is NP-hard, so the search itself scales poorly as models and device counts grow.
Shard the optimizer state (ZeRO). Attack the replication problem by splitting optimizer states, gradients, and parameters across devices instead of copying them. Where it broke. It genuinely saves memory, but pays for it with extra collectives — reduce-scatter and all-gather — so it trades a memory cost for more communication, still within a spatial framing.
Pipeline parallelism. Cut the model into stages of layers on different GPUs and stream micro-batches through them, using only cheap point-to-point messages. Where it broke. It sidesteps collectives but introduces pipeline bubbles from periodic flushes, and it is orthogonal to the tensor-partition costs rather than a cure for them.
The gap. Every approach kept each sub-operator pinned to its own device, so the whole design space was spatial, and both all-reduce and replication were laws of that space rather than choices within it. What was missing was a way to place two sub-operators on the same device at different times — a temporal axis that could let a device finish a sum by itself instead of over the network.

How PrimePar works

The idea

Add a time axis so a device can finish its own sum, and the expensive collective disappears. The whole cost structure of spatial partitioning comes from one rule that nobody had questioned: every piece of a split operator lives on its own device. Under that rule, splitting a summed-over dimension is doomed to scatter partial sums that only an across-the-network all-reduce can gather. PrimePar breaks the rule by adding a temporal dimension — letting two pieces run on the same device at different times. Now the summed-over dimension can be walked through step by step on one device, which simply accumulates the running total by itself, and the all-reduce that consumed a third to two-thirds of training never happens. The leftover coordination is reduced to small ring messages between neighbors that overlap with the compute of each step. Restructuring a parallel-training layout so that an expensive, stall-inducing bottleneck is designed out rather than merely sped up is the same move the DNN-training-pipeline work covered elsewhere on this site makes when it reshapes how a model is split across devices so the costly stage stops stalling the rest.

Move data instead of copying it, and the memory waste goes too. The second cost of pure spatial partitioning is replication: any tensor whose dimensions a plan does not cut has to be copied onto every device that needs it, and that waste grows steadily worse as parallelism scales. The temporal primitive inherits Cannon’s systolic pattern, in which tiles shift around the grid from one neighbor to the next rather than being broadcast, so a shared tensor exists in a single place and travels to where it is next needed. Reuse replaces duplication, and peak memory drops to as little as 68 percent of the spatial baseline on the largest models. Cutting the redundant memory traffic of a data-parallel workload by reusing data already close at hand instead of re-fetching or re-copying it is the same economy the GPU graph-learning work covered elsewhere on this site pursues, when it reuses the neighbor data that overlapping subgraphs share rather than loading it again for each.

Make the bigger space searchable by exploiting the model’s repetition. Adding a temporal axis multiplies an already exponential space of partition plans, and the prior integer-programming search was NP-hard and scaled badly. PrimePar makes the search tractable with a segmented dynamic program: the transformer’s graph has awkward skip connections that break a naive recurrence, so the method cuts the graph into segments where a clean recurrence holds, solves each, and merges them. The decisive economy is that a transformer is a stack of identical layers, so the optimal plan for one layer is computed once and merged across layers in a logarithmic number of steps, letting a 175-billion-parameter model be planned in seconds on one CPU thread. Turning an intractable global search into a set of small local sub-problems whose solutions compose back into the whole is the same discipline the scalable hardware-invariant-learning work covered elsewhere on this site relies on, decomposing a proof that could never be checked in one piece into many small pieces that combine into the full guarantee.

Following the reasoning through PrimePar

An MLP block on 8 GPUs, Megatron versus PrimePar

The paper's ablation traces the same OPT-175B MLP block through both systems on 8 GPUs, making the communication difference concrete.

Step 1 of 5

Megatron splits a summed-over dimension spatially

For the second linear layer of the MLP block, Megatron-LM partitions the summed-over hidden dimension across devices. Each device computes a partial sum of the output, so the block ends with an all-reduce of the output tensor within each node of four GPUs.

This intra-node all-reduce sits on the critical path; across the block, collective communication is 34 to 62 percent of the MLP latency.

Why it works

Accumulating over time eliminates the reduction instead of accelerating it. Splitting a summed-over dimension across devices leaves each with a partial sum that only all-reduce can combine → but sequencing those same pieces across time steps on one device lets it add them into its own output locally → so the collective that ate 38 to 65 percent of training simply never happens, rather than merely running faster.
Shifting tiles in a ring reuses shared data instead of copying it. A purely spatial split replicates any tensor whose dimensions it does not cut, wasting memory on every device → but the Cannon-style schedule keeps each tile in one place and passes it neighbor to neighbor → so a shared tensor is reused across steps rather than duplicated, and peak memory falls to as little as 68 percent of the spatial baseline.
The only remaining traffic is point-to-point, which hides behind computation. All-reduce sits on the critical path and stalls the next step until the global sum is known everywhere → but ring messages travel just to an adjacent device and run on a separate stream → so they overlap with each step’s matrix multiply and add almost nothing to wall-clock time.
It changes the communication pattern without touching compute efficiency. Many parallelization tricks buy cheaper communication by shrinking tiles or lowering arithmetic intensity → but the temporal primitive keeps each device’s matrix multiplies the same size as before → so its computation latency matches Megatron-LM step for step, and the speedup is pure communication savings.
Exploiting the transformer’s repeated layers makes searching the bigger space tractable. Adding a time axis multiplies an already exponential partition space, and the prior integer-programming search is NP-hard → but a segmented dynamic program solves one layer’s optimal plan and merges identical layers in a logarithmic number of steps → so the best strategy for a 175-billion-parameter model is found in seconds on a single CPU thread.

What it costs (trade-offs and scope)

The speedups are modest for smaller models. Throughput gain is 1.16 to 1.20 times for models around 7 billion parameters; the larger 1.11 to 1.68 times figures and the 1.30-times geometric mean come from models above 100 billion at bigger scale. A small model on a small cluster sees much less.
The memory win is scale-dependent too. PrimePar uses about 90 percent of Megatron-LM’s memory for 7-billion models; the headline 68 percent appears only for the largest models, BLOOM 176B, at 16 to 32 devices.
The primitive needs a structured, power-of-two device grid. The square-grid primitive arranges a power-of-two count of homogeneous devices as a square and gives each an equal run of time steps. It is clean on regular clusters but assumes that regularity, rather than adapting to arbitrary or heterogeneous device counts.
The overlap only holds when the interconnect keeps up. Ring communication hides behind computation only while each step’s ring message is no longer than its matrix multiply; on an interconnect that is slow relative to compute, the overlap breaks down. Torus topologies like TPU v4 are argued to be the ideal fit, but that is projected, not measured here.
Three-phase alignment is a hard design constraint. The forward, backward, and gradient phases must keep their tensor distributions lined up at phase boundaries, or the system pays for redistribution; the primitive is carefully constructed to satisfy this, which constrains how freely it can be composed.
The search is a heuristic decomposition. The segmented dynamic program rests on structural assumptions about the transformer graph and on its stack of identical layers; it is exact under those assumptions, not a general optimizer for arbitrary computation graphs.
It improves the tensor-partition layer, not the whole stack. PrimePar enhances data and model parallelism and composes into 3D parallelism, where it adds up to 1.46 times, but still relies on pipeline parallelism and standard frameworks; it was evaluated on 32 V100 GPUs, with larger-cluster behavior argued rather than demonstrated.

The evaluation

Ways to parallelize a transformer operator across devices

Property	Megatron-LM (manual)	ZeRO (shard state)
Partition space
Summed-over split needs all-reduce
Shared tensors in memory
How the plan is found
Result relative to Megatron-LM	Baseline	Less memory, extra comm

Parallelization approaches for transformer training, evaluated on OPT, Llama2, and BLOOM from 7 billion to 176 billion parameters across up to 32 V100 GPUs. PrimePar cuts collective communication to 19.9 to 62.2 percent of Megatron-LM's, reaches up to 1.68 times throughput at 69 percent peak memory, a 1.30-times geometric-mean speedup at 32 GPUs, and up to 1.46 times when folded into 3D parallelism, with the plan found in seconds.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

PrimePar is a genuinely elegant idea, cleanly executed: by adding a temporal axis to tensor partitioning, it lets a device finish a summed-over reduction locally over time instead of over the network, designing out the all-reduce and the replication that were baked into every prior spatial-only system, and it finds the plans automatically in seconds. Read the numbers with their scope. The gains scale with the problem: only 1.16 to 1.20 times and about 90 percent memory at 7 billion parameters, with the striking 1.68 times and 68 percent figures reserved for models above 100 billion at larger device counts. The primitive assumes a structured, power-of-two square grid of homogeneous devices, and its communication only fully hides behind computation when the interconnect is fast relative to compute — the ideal torus case is argued, not measured. It carries a real design constraint in keeping the three training phases aligned, and the search, though fast, is exact only under structural assumptions about the transformer graph. It was demonstrated on 32 V100 GPUs, with very-large-cluster behavior projected. The durable contribution is the reframing itself: a whole dimension of parallelism that the field had left on the table.

Background

Why training the largest models is a communication problem

It helps to see why so much effort goes into how tensors are laid out across GPUs rather than into the math itself. A modern large language model has more parameters and activations than any single accelerator can hold, so training it is inherently a distributed computation, and the arithmetic is embarrassingly parallel. What is not parallel is the coordination: the moment you split a computation whose parts must be summed, the devices have to talk, and at the scale of hundreds of billions of parameters that talking — the collectives, the replication, the synchronization — comes to dominate the wall-clock time and the memory bill far more than the floating-point work does. So the real engineering frontier for large-model training is not faster multipliers but smarter placement, choosing how to cut and schedule the tensors so the machines spend their time computing rather than waiting on each other. PrimePar answers that by inventing a placement dimension nobody was using. The same conviction that training performance is won or lost across the whole stack drives the Trainium accelerator work covered elsewhere on this site, which argues that fast large-model training is a full-stack co-design problem spanning silicon, interconnect, and software rather than any one layer. The broader pattern is that as models outgrow every device, the systems that win are the ones that treat communication, not computation, as the thing to minimize.

If you remember only three things

1. Large transformers are trained by tensor partitioning, splitting each operator across GPUs, but every prior system distributes the pieces only across space — one per device — so splitting a summed-over dimension leaves each device with a partial sum that must be gathered by all-reduce (38 to 65 percent of training time and hard to overlap), and unpartitioned tensors get replicated on every device, wasting memory that worsens with scale.
2. PrimePar adds a temporal dimension: borrowing Cannon's systolic matrix-multiply, it lets a device run a short sequence of steps in which tiles shift around a logical grid and partial sums accumulate locally, so the reduction finishes in place with no all-reduce, shared tensors are reused rather than replicated, and the only leftover traffic is cheap ring messages that overlap with computation; a segmented dynamic program exploits the transformer's identical layers to search this larger space in seconds.
3. On up to 32 V100 GPUs across OPT, Llama2, and BLOOM, PrimePar reaches up to 1.68 times throughput at 69 percent peak memory, a 1.30-times geometric-mean speedup at 32 GPUs, and up to 1.46 times folded into 3D parallelism, with collective communication cut to a fifth-to-two-thirds and compute unchanged — but the gains are modest for small models (1.16 to 1.20 times at 7 billion), the memory win is largest only at the biggest scale, the primitive assumes a power-of-two square grid, and the overlap holds only when the interconnect keeps pace with compute.

Check yourself

Check your understanding

1. Why does purely spatial tensor partitioning force expensive all-reduce communication?
2. What is PrimePar's core new idea?
3. How does PrimePar make searching its larger partition space fast enough to be practical?
4. Which statement best reflects the limits of PrimePar's results?

Sources

The paper is “PrimePar: Efficient Spatial-temporal Tensor Partitioning for Large Transformer Model Training” by Haoran Wang, Lei Wang, Haobo Xu, Ying Wang, Yuming Li, and Yinhe Han, published at ASPLOS 2024, pages 801 to 812, DOI 10.1145/3620666.3651357. The authors are at the Institute of Computing Technology, Chinese Academy of Sciences, and the University of Chinese Academy of Sciences.

This explainer is written from the full 17-page text. Everything above comes from the paper’s own sections: the introduction and motivation showing that spatial-only tensor partitioning forces all-reduce on summed-over splits (38 to 65 percent of Megatron-LM training latency) and replicates unpartitioned tensors, with the temporal-dimension insight drawn from Cannon’s parallel matrix multiplication; the partition-space section formalizing sub-operators by device identifier and temporal step through Dimension Slice Indices, the existing partition-by-dimension baseline, and the novel square-grid primitive that is collective-communication-free, memory-efficient, and training-supporting via forward, backward, and gradient phase alignment; the cost model of intra-operator communication, computation, and peak memory with communication latencies fit by profiling and linear regression and computation overlapped with ring communication; the segmented dynamic-programming optimization that decomposes the transformer graph and exploits its identical stacked layers to plan in seconds; and the evaluation on a 32-GPU cluster of V100s across OPT, Llama2, and BLOOM from 7 billion to 176 billion parameters, reporting the 1.16-to-1.68-times throughput gains, the 1.30-times geometric-mean speedup at 32 GPUs, peak memory down to 68 percent, collective-communication latency cut to 19.9 to 62.2 percent of Megatron-LM’s with unchanged computation latency, the up-to-1.46-times improvement folded into 3D parallelism, and optimization times of seconds even at 32-way parallelism.