Breaking the computation and communication abstraction barrier in distributed machine learning workloads

Machine-learning frameworks call one library for math and another for network transfers, so neither can optimize across the boundary — CoCoNet expresses both in one small program, then mechanically splits, reorders, fuses, and overlaps them, turning roughly 2,000 lines of hand-written CUDA into 15 lines of DSL and speeding up BERT, GPT-2, and GPT-3 workloads by 1.2–1.8×.

Why this problem is worth solving

Models outgrew single machines years ago — BERT at 340 million parameters, GPT-2 at 1.5 billion, GPT-3 at 175 billion — so every serious training or inference job is distributed, and its inner loop interleaves two kinds of work. Computation: tensor algebra on each GPU, served by libraries like cuBLAS and cuDNN. Communication: exchanging tensors between GPUs, served by collective libraries — NVIDIA’s

The two are separate abstractions: a framework calls a computation kernel, waits, calls a communication kernel, waits. The paper’s opening measurement shows what the wall costs. In model-parallel inference of GPT-2, a matrix multiply is followed by an AllReduce of its output. Run sequentially — as every framework does — the network sits idle during the multiply and the GPUs sit largely idle during the reduce. Overlapping the two in a fine-grained way hides 80% of the matrix multiply’s execution time, a 1.36× speedup on 16 V100 GPUs — but the hand-written implementation of that one overlap is about 2,000 lines of CUDA, valid for one operator pair, one topology, one library version.

The paper’s diagnosis: the barrier is an artifact of software packaging, not of hardware. Its bet: a compiler can cross it mechanically.

What people did before, and where it broke

Independent optimization inside each box. Kernel fusion is standard for computation — ML compilers merge chains of element-wise operations into one kernel (see the TensorIR and AStitch explainers elsewhere on this site). NCCL tunes its collectives with topology-aware protocols and channels. Neither touches the other: the output of an AllReduce is always written to memory, then re-read by the next computation kernel.

Fixed cross-boundary strategies baked into systems. ZeRO and GShard distribute optimizer state and computation across ranks — effectively one specific reorder of communication and computation, hard-coded into a runtime. Valuable, but a point in the design space, not a way to explore it: the paper shows these schedules lose to alternatives at small tensor sizes and leave fusion and overlap opportunities unexploited.

Coarse-grained overlap. Frameworks overlap a layer’s gradient AllReduce with the previous layer’s backward pass — scheduling whole kernels concurrently. Fine-grained overlap inside a producer-consumer pair (start communicating a chunk the moment it is computed) requires controlling the order in which the producer kernel emits its output, which no library interface exposes.

Manual heroics. Where the payoff justified it, engineers wrote custom fused kernels — thousands of lines each, unmaintainable across the matrix of models, sizes, and clusters.

The idea

Three moves make the barrier crossable.

Give tensors a distribution type. In CoCoNet a tensor is not just shape and dtype but also layout: sliced (each of N ranks holds 1/N along some dimension), replicated (every rank holds the same full value), or local (per-rank values, same shape). Operations infer their output layouts — a MatMul of a sliced weight produces a local partial result; an AllReduce turns local into replicated; a ReduceScatter turns local into sliced. This bookkeeping, absent from every existing framework, is what lets the compiler prove that a rewrite across the boundary preserves meaning.

Separate algorithm from schedule — Halide’s lesson, imported. The program says what is computed; a schedule — an ordering of four transformations — says how:

• split breaks one collective into finer ones. The key instance splits AllReduce into a
  sum across ranks where each rank keeps only its 1/N slice of the result
  followed by an
  concatenation of all slices so every rank ends with the full tensor
  . Always valid — and it exposes a seam.
• reorder slides computation into that seam: operations after the AllGather move before it, becoming sliced computations on 1/N of the data. Valid exactly when the operations can be sliced along the gathered dimension. This one transformation reproduces ZeRO- and GShard-style distribution as a special case.
• fuse merges a ReduceScatter, the sliced computations, and the AllGather into a single FusedAllReduce kernel: values arriving from the network flow through registers into the computation instead of taking a store-and-reload round trip through GPU memory.
• overlap runs producer and consumer concurrently at chunk granularity — the matrix multiply is scheduled to produce output chunks in precisely the order the ring AllReduce transmits them, so communication of chunk k rides under computation of chunk k+1, in one kernel launch each.

Let a machine pick, because no schedule always wins. The paper’s own data kills the idea of a universal best: plain AllReduce-then-fused- computation wins for tensors up to about 2¹⁶ elements, while the fully fused-and-distributed schedule wins beyond 2¹⁷ — fusion can even lose by inflating register usage and cutting GPU occupancy. The autotuner turns that judgment call into a ten-second search.

MM → AR → (bias + Dropout + residual)
every GPU: pointwise work on N/N of the data
AR output: stored, then reloaded

Why it works

• Tensors carry distribution layouts → the compiler can statically verify that a rewrite across the computation/communication boundary preserves semantics → transformations that were one-off engineering become reusable, composable, and safe.
• AllReduce is exactly ReduceScatter plus AllGather → splitting it creates a seam where computation can run on 1/N of the data → work and (for the optimizer) state distribute across ranks for free.
• Fusing communication with computation keeps intermediates in registers → the memory round trip between kernels disappears → bandwidth-bound pointwise stages accelerate and kernel-launch overhead drops.
• The producer’s output order is compiler-controlled → chunks are computed in the ring’s transmission order → communication hides under computation at fine grain, using compute and network (and NVLink and InfiniBand simultaneously, in the pipeline case) instead of alternating them.
• The best combination depends on size and topology → an autotuner that executes real candidates picks it in seconds → users get hand-tuned performance from 15-line programs, on every configuration.

What it costs

• A new language to adopt. Workloads must be expressed in the DSL (or wrapped by someone who does); the paper’s programs are 12–20 lines, but they sit outside the everyday PyTorch idiom.
• NCCL surgery. Fused collectives are built by modifying NCCL’s internal kernels — powerful, but coupled to one communication library’s architecture and evolution.
• Autotuning per configuration. Schedules are chosen by measuring real runs for each hardware, topology, and size; cheap (seconds) but repeated whenever the setup changes.
• Fusion can backfire. Fused kernels raise register usage and can lower occupancy — the paper’s own small-tensor results show the fused schedule losing, which is precisely why the autotuner exists.
• Scope. Transformations center on collectives with clean algebraic structure (AllReduce, AllGather, point-to-point); the evaluation is NVIDIA-only and inference/training loops of transformer workloads.

The evaluation

All experiments ran on 16 NVIDIA DGX-2 nodes — 256 Tesla V100 GPUs total, NVSwitch/NVLink inside nodes and InfiniBand between them. CoCoNet was applied to all three parallelism styles, each against the strongest available hand-optimized baseline.

Two details deserve emphasis. First, the interface optimization — letting kernels take a list of scattered per-layer buffers instead of one contiguous tensor — is mundane but real: NVIDIA’s own BERT script spends time each iteration copying every layer’s gradients into a single buffer just to satisfy AllReduce’s calling convention, then copying results back. Second, the memory effect of the reorder transformation is a capability change, not just a speedup: distributing Adam’s momentum and variance state let a 3.9-billion-parameter BERT train with plain data parallelism, which NV BERT and PyTorch DDP simply cannot do.