Centauri: Enabling Efficient Scheduling for Communication-Computation Overlap in Large Model Training via Communication Partitioning

Hiding communication behind computation in LLM training has been tried with fused kernels and with graph schedulers, each missing what the other sees — Centauri observes that any collective can be partitioned three ways (substitute the primitive, split the device group, chunk the workload) and schedules the pieces at operation, layer, and model tiers, for up to 1.49× faster training.

Why this problem is worth solving

This page completes a trilogy on this site. CoCoNet (ASPLOS 2022) broke the computation/communication abstraction barrier with a DSL and fused GPU kernels; Google’s decomposition work (ASPLOS 2023) made overlap automatic in XLA for TPU pods. Centauri arrives at the GPU-cluster frontier where LLM training actually happens — hybrid parallelism on heterogeneous networks — and argues both prior philosophies explore corners of a larger space.

The setting: pre-training scales by combining

The paper’s survey table makes the fragmentation concrete: each prior system — Megatron-LM, DeepSpeed ZeRO, PyTorch FSDP, out-of-order backprop, breadth-first pipelines, CoCoNet, Dist-Einsum — exercises a subset of the possible partitioning and scheduling dimensions. Nobody had named the whole space, let alone searched it.

What people did before, and where it broke

Graph-level scheduling. FSDP prefetches the next layer’s weight allgather during the current layer’s compute; ZeRO overlaps gradient reduction with backward. Effective but coarse: whole collectives overlap whole computations, and the paper’s Figure-1 timeline shows the tell-tale idle gaps when a collective outlasts its computation partner.

Multi-stream concurrency. Running collectives on parallel GPU streams (Better-Together-style) parallelizes communication with communication, not with computation — a narrower win.

Fine-grained kernel fusion. CoCoNet and the decomposition approach split a collective and its dependent matmul into chunks inside fused kernels — powerful for that pair, but locally greedy: Centauri’s motivating example shows the unfused MatMul B finishing earlier under a better graph-level order than the fused kernel achieves. Fusing an operator pair forecloses reorderings around it.

Single-parallelism focus. Most schedulers optimize one parallelism’s pattern. Hybrid reality is messier: the best backward schedule under tensor parallelism (activation gradients first) differs from the best under FSDP (weight gradients first) — a choice prior systems hard-coded one way.

The idea

Name the space first. Centauri’s conceptual move is the observation that a collective is nothing but a mapping transformation of a data volume across a device group — so there is nothing sacred about its granularity. It can be rewritten (primitive), subdivided across machines (group), or sliced across data (workload), and any combination preserves semantics. The three dimensions are searched in that order because each stage’s choice sets the cost model for the next. Two satisfying corollaries: the optimization spaces of eight prior systems are subsets of this space, and famous parallelism strategies themselves — ZeRO-1/2, sequence parallelism — turn out to be specific substitution-and- repositioning choices within it. What the field treats as distinct algorithms, Centauri treats as coordinates.

Prune with topology, not hope. The space is huge, and much of it is bad. Pruning rules encode network physics: rooted primitives lose to unrooted ones as device counts grow (root contention); sub-groups mixing fast and slow links are excluded (the slow link sets the pace); chunk counts are capped (launch overhead); on some architectures finer group splits would fragment NCCL’s ring channels and are rejected. What survives is a tractable menu of genuinely plausible partitions.

Schedule where the structure is. A global schedule over a partitioned billion-parameter graph is intractable — but LLMs are towers of identical layers trained with gradient accumulation, so the problem factorizes. Operation tier: within a forward layer, each collective greedily picks its best partition-and-overlap scheme (Algorithm 1 sizes chunks by the communication-to- independent-compute ratio and orders intra/inter-node pieces). Layer tier: the backward pass has a natural binary choice — activation gradients first (the classical order, right for TP) or weight gradients first (right for FSDP, whose next-layer weight allgather wants to start early) — decided per configuration by comparing two cost formulas. Model tier: the DP gradient allreduce is split into an allgather hidden under the next forward pass and a reduce-scatter hidden under backward, chunk-by-chunk; and in pipelines, the number of micro-batches launched together interpolates between depth-first (minimum activation memory, minimum overlap) and breadth-first (maximum both) — chosen to maximize overlap within the memory budget.

Different clusters, different winners — automatically. Because every choice is costed against the actual topology, the same model gets different schedules on the bandwidth-starved Cluster A (25 Gb/s/device inter-node) and the well-fed Cluster B (100 Gb/s) — and the gains concentrate exactly where bandwidth is scarce, which is where they are needed.

AR → RS + AG (phase-splittable)
AG → [intra-node × inter-node] × chunks

Why it works

• A collective is a workload mapping over a group → partitioning by primitive, group, and workload preserves semantics while multiplying overlap opportunities → the spaces of eight prior systems (and even ZeRO/SP themselves) become searchable points rather than competing doctrines.
• Topology-aware pruning encodes network physics → unbalanced sub-groups, rooted contention, and over-fine chunks are excluded a priori → the surviving space is small enough for greedy, per-collective decisions.
• LLM structure factorizes the schedule → identical layers and gradient accumulation let three cheap tiers (operation, layer, model) compose into a global schedule → no monolithic graph-scheduling problem is ever solved.
• The backward pass has a genuine binary degree of freedom (activation- vs weight-gradient priority) → costing it per configuration instead of hard-coding captures wins that single-parallelism schedulers structurally miss.
• Overlap potential is bounded by concurrent computation → the micro-batch count knob converts spare activation memory directly into hidden communication → memory, not cleverness, becomes the binding constraint, and it is spent deliberately.

What it costs

• Search and modeling machinery. Cost models for every partition scheme, per topology; greedy tier-wise choices are cheap but the framework must be maintained against evolving NCCL/cluster behavior.
• Memory spent on overlap. The model-tier knob deliberately raises activation memory (more concurrent micro-batches) to buy overlap — capacity that competes with model size and batch size.
• Kernel-launch granularity limits. Chunking is capped by launch overheads; on very fast networks the whole approach yields little (~5% at 256 GPUs on Cluster B) because there is little to hide.
• Greedy, not optimal. Operation-tier decisions are sequential and local; the hierarchy trades global optimality for tractability.
• Training-loop scope. The techniques lean on LLM regularity (uniform layers, gradient accumulation); irregular models or inference serving are outside the demonstrated envelope.

The evaluation

LLaMA models at four sizes trained under FSDP, DP, FSDP+DP, and TP+PP+DP on two deliberately different clusters: Cluster A (8× A100-80G per node, 600 GB/s NVSwitch intra-node, but PCIe 3.0 hosts and a shared InfiniBand HDR port — 25 Gb/s per device inter-node) and Cluster B (4× HDR ports — 100 Gb/s per device), scaling to 256 GPUs. Baselines: Megatron-LM, DeepSpeed ZeRO-1/2/3 with direct overlap, and per-configuration overlap variants.