Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads

1 INTRODUCTION

• Breaking the abstraction boundary enables interface, fusion, reordering, and fine-grained overlap optimizations that separate computation and communication kernels cannot provide.For example, accepting a list of non-contiguous parameter buffers can avoid copying them into a single buffer before AllReduce.Source: §1 · pp. 1–2
• Fine-grained MatMul–AllReduce overlap hides 80% of MatMul execution time and provides 1.36× speedup by communicating each output chunk as soon as it is computed.The MatMul kernel computes chunks in the ring AllReduce communication order, while one MatMul and one AllReduce kernel avoid multiple launch overheads.Source: §1 · p. 2, Figure 1 · p. 2
• Manually implementing scenario-specific optimizations is laborious; the described overlapping implementation alone contains approximately 2,000 lines of CUDA code.As models grow, training and inference must be distributed, and even small performance gains can save time, energy, and money [49].Source: §1 · pp. 1–2

transformations.

• CoCoNet combines a DSL for computation and communication with scheduling, autotuning, and code generation to produce optimized kernels for specific hardware, topology, and data sizes.Its autotuner applies transformations across the computation–communication boundary.Source: §1 · p. 2, transformations · p. 2
• CoCoNeT speeds BERT training by up to 1.68×, model-parallel inference by up to 1.51×, and pipeline-parallel inference by up to 1.77× across the reported models.The training result covers BERT models; model-parallel inference covers BERT 3.9-billion- and GPT-2 8.2-billion-parameter models; pipeline-parallel inference covers GPT-2 8.2-billion- and GPT-3 175-billion-parameter models.Source: transformations · p. 2

2.1 Tensor Layout

• CoCoNet tensors specify datatype, shape, and distribution layout, which may be sliced across ranks, replicated with identical values, or local with rank-specific values.A scalar is a zero-dimensional tensor available on all ranks.Source: §2.1 · p. 2

2.2 CoCoNet’s Operations

• CoCoNet programs use data-flow graphs with local computation and cross-rank communication operations, while operation rules infer intermediate tensor shapes and layouts.For example, MatMul on two sliced tensors yields a local tensor, and AllReduce returns a replicated tensor; Execute names the program, inputs, and outputs.Source: §2.2 · p. 3

2.3 Fused Collective Communication Operations

• FusedAllReduce passes communication output directly to following computations through registers, avoiding the intermediate memory stores and reloads required by separate operations.It accepts computations as additional arguments to AllReduce.Source: §2.3 · p. 3

2.4 Overlapping Operations

• The Overlap construct lets dependent operations such as MatMul and AllReduce run in an overlapped manner to use network and computation resources.The paper identifies the consecutive MatMul and AllReduce in Figure 3 as an example.Source: §2.4 · pp. 3–4

2.5 Custom Operations

• Custom operators define syntax, output shape and layout semantics, and code generation, which can call existing libraries or generate fused GPU kernels.Syntax and semantics can take only a few lines, while code generation for complex operations may take hundreds; optimized library implementations can be called instead.Source: §2.5 · p. 4

3.1 Splitting Communication

• AllReduce Split RS-AG replaces an AllReduce with a ReduceScatter that produces a sliced tensor and an AllGather that returns a replicated tensor; the split is always valid.The user can apply this split as part of a manually specified schedule or an autotuned schedule.Source: §3.1 · p. 4, §3 · p. 4

3.2 Reordering Operations

• AllGather Reorder moves an AllGather across communication or computation by replacing computations with equivalent sliced operations and gathering their outputs.It is valid only when the reordered operations can be sliced along the AllGather dimension; slicing rules depend on operation type and input dimensions.Source: §3.2 · p. 4

3.3 Fusing Operations

• CoCoNet fuses computation sequences or combines ReduceScatter, sliced computations, and AllGather into a FusedAllReduce, provided the fusion preserves DFG dependencies.In the running example, ReduceScatter, two sliced computations, and AllGather are fused into one FusedAllReduce.Source: §3.3 · p. 4

3.4 Overlapping Operations

• The overlap transformation runs producer-consumer operations on multiple hardware resources simultaneously, and is valid only when the operations have producer-consumer relationships.The running example overlaps a matrix multiplication with a FusedAllReduce.Source: §3.4 · p. 5

3.5 Automatic Exploration of Schedules

• CoCoNet’s autotuner searches schedules for the best performance on the target architecture and input sizes, returning the minimum-time schedule after breadth-first exploration.It first fuses pointwise computations up to a threshold to shrink the search space; the passage reports exploration takes only a few seconds for all workloads.Source: §3.5 · p. 5

4 DISTRIBUTED WORKLOADS IN COCONET

• CoCoNeT fuses Adam updates into one kernel, splits AllReduce into ReduceScatter and AllGather, reorders computation, and slices optimizer states to remove their AllGather.Figure 6 covers parameters (p), gradients (g), momentum (m), and velocity (v); the baseline averages gradients with AllReduce before updating state and parameters.Source: §4 · p. 5, Figure 6 · p. 5

GPU (0,1)

• Splitting AllReduce into ReduceScatter and AllGather and reordering P2P sends with AllGather avoids sending replicated output, reducing inter-group communication by group size.Source: GPU (0,1) · p. 6, Figure 7 · p. 6
• The optimized pipeline schedule fuses P2P sends with computation and overlaps the resulting operations; tiled buffers allow intra- and inter-group communication to overlap.Source: GPU (0,1) · p. 6

5 THE COCONET CODE GENERATOR

• CoCoNet generates collective calls, fused-computation kernels, fused-collective kernels, and kernels that overlap communication with computation.It traverses the program’s dataflow graph to generate kernel calls, wraps generated programs as custom operators, and supports operations on multiple non-contiguous tensors.Source: §5 · p. 6

5.1 NCCL Architecture

• CoCoNet adapts NCCL’s configuration of topology, protocols, channels, and threads to custom communication and computation, using NCCL’s tiled communication workflow.NCCL divides collective data into buffer tiles, then chunks across ranks and channels; channel threads copy data and apply reductions when needed.Source: §5.1 · pp. 6–7

5.2 Fused Collective Communications

• CoCoNet extends NCCL collectives to fuse arbitrary pointwise computations and reductions, passing intermediate values directly through registers.It generates code for all NCCL protocols and element types, dispatching at runtime; it also handles mixed element types, sliced-tensor indexing, and reductions on slices using existing connections.Source: §5.2 · p. 7

5.3 Overlapping of Communication and Computation

• Up to 1.36× better performance and more than 80% of MatMul time hidden through fine-grained MatMul–AllReduce overlap.CoCoNet uses chunk-ordered computation and communication with fine-grained synchronization, while retaining full-matrix MatMul execution and communication bandwidth.Source: §5.3 · p. 7, Figure 9 · p. 8

5.4 Operations on Scattered Tensors

• Scattered-tensor overhead is insignificant over a contiguous tensor; for BERT with 334M elements, bucket metadata requires 0.6% extra memory.CoCoNet assigns tensor buckets to warps for efficient indexing, precomputing offsets; bucketing occurs once on the CPU and its computation overhead is negligible across thousands of training iterations.Source: §5.4 · pp. 7–8, Table 2 · p. 8

5.5 PyTorch Integration

• CoCoNet operations are exposed through PyTorch’s torch.distributed module, reusing NCCL initialization logic and enabling models to invoke them through wrapper functions.The wrappers prepare operation arguments, including scattered-tensor bucket pointers and cleared spin-lock buffers for overlapping.Source: §5.5 · p. 8

6.1.1 Standalone Experiments.

• In data-parallel standalone experiments, CoCoNet provides 1.2× to 1.7× speedup for Adam and 1.35× to 2.0× for LAMB over baselines, with the best schedule varying by tensor size.AR-Opt performs well on smaller tensors, while fused RS-Opt-AG approaches optimal speedups on larger tensors; fusion can hurt smaller tensors because of higher register usage, and autotuning explored schedules in 10 seconds.Source: Figure 10 · p. 9, §6.1.1 · p. 9

6.2.1 Standalone Experiments.

• For GPT-2 8.3B model-parallel operations, CoCoNet’s best schedule provides 1.42×–1.70× speedup over Megatron-LM and 1.21×–1.34× over GShard-Eq.The schedule overlaps FusedAllReduce with matrix multiplication; the autotuner explored schedules in 12 seconds.Source: Figure 11 · p. 10, §6.2.1 · p. 10

6.1.2 Integeration with BERT.

• For BERT 1.2B with Adam, CoCoNet provides 1.53× speedup over NV BERT and PyTorchDDP; for LAMB, it provides up to 1.64× speedup over all baselines.On BERT 3.9B, NV BERT and PyTorch ran out of memory; distributing optimizer state reduces memory use and can enable larger batch sizes.Source: §6.1.2 · p. 10, Table 4 · p. 11

6.3.1 Standalone Experiments.

• For GPT-3 175B pipeline parallelism, CoCoNet is 11.75×–12.21× faster than Megatron-LM, 2.84× faster than AR-C-P2P-AG, and 1.66×–1.72× faster than GShard.Sliced P2P reduces cross-node communication volume; fusing operations improves memory-bandwidth utilization, and overlapping communication across NVLink and InfiniBand improves network-bandwidth utilization.Source: Figure 12 · p. 11, §6.3.1 · p. 11

6.3.2 Integration with Megatron-LM.

• Integrating CoCoNet’s pipeline-parallel schedule in Megatron-LM significantly improves inference throughput for GPT-2 and GPT-3 through fusion and fine-grained overlap of communication operations.The integrated schedule is ol(RS,fuse(C-P2P),AG).Source: §6.3.2 · p. 11

Distributed Machine Learning Abstractions

• Unlike existing distributed machine-learning abstractions, CoCoNeT preserves tensor layouts for static type checking and automatic program transformations.Prior frameworks [13] [29] [40] [45] and DSLs [18] express tensor-in/tensor-out programs but do not provide these layout-enabled capabilities.Source: §7 · pp. 11–12

Distributed Neural Network Training

• CoCoNeT adds scattered tensors, fused collectives, and communication-computation overlap beyond prior distributed training techniques.Horovod [45] fuses gradients in a 64MB buffer, while PyTorch DDP [35] overlaps AllReduce but requires extra memory and lacks optimizer-update slicing. CoCoNeT also overlaps inter- and intra-node communication; BytePS is complementary.Source: §7 · p. 12

Optimizing Stencil Computations

• CoCoNeT extends stencil DSLs to reason about and compose collective communication with computation for distributed machine learning.Prior systems include algorithm-schedule separation in Halide and Fireiron, kernel generation in TVM [18], and distributed scheduling in Distributed-Halide.Source: §7 · p. 12

Overlapping Computation and Communication

• Unlike prior overlap techniques, CoCoNeT overlaps multiple GPU computations and communication operations, including collectives with convolutions and matrix multiplications.Prior approaches use pipelining or non-blocking MPI; cited MPI+OpenMP techniques target CPU stencil halo exchanges rather than these GPU workloads.Source: §7 · p. 12

8 CONCLUSION

• CoCoNeT is a language for describing distributed machine learning workloads and optimizing them across the computation–communication boundary.The paper presents it as a way to address both computation and communication together.Source: §8 · p. 12
• CoCoNeT-generated code significantly improves several large language model training and inference times.Source: §8 · p. 12
• The authors plan to automate the optimizations through smart search in future work.Source: §8 · p. 12

• The artifact contains CoCoNet’s implementation source code and benchmarking infrastructure to reproduce all Section 6 results.Source: §9 · p. 12

A.1 Abstract

• The appendix explains how to reproduce standalone experiments in Figures 10–12 and integration results in Sections 6.1.2, 6.2.2, and 6.3.2.These are the experiments covered by the artifact’s reproduction instructions.Source: §A.1 · p. 12
• The artifact includes the CoCoNeT DSL and compiler, plus generated code integrated with PyTorch, Megatron-LM, and NVIDIA Bert.Source: §A.1 · p. 12
• Reproduction should use a system similar to the authors’ experimental system.Source: §A.1 · p. 12
• All experiments can also run on a system with more than one NVIDIA GPU.Source: §A.1 · p. 12

A.2 Artifact Check-list (Meta-information) •

• The artifact uses a C++ CoCoNet DSL/compiler to generate MPI programs with CUDA kernels; evaluation used 256 Tesla V100 GPUs across 16 DGX-2 nodes.The checklist also lists BERT, GPT-2, and GPT-3 datasets, about 100 GB of disk space, and roughly one hour for preparation and five hours for experiments.Source: §A.2 · pp. 12–13

A.3.1 How to Access.

• The CoCoNet implementation and benchmarking infrastructure are publicly available as artifact, in separate coconet and coconet-experiments directories.The passages provide GitHub locations for the latest versions of both directories.Source: §A.3.1 · p. 13

A.3.2 Hardware Dependencies.

• Benchmarks run on systems with two or more NVIDIA GPUs, while reproducibility is specified for the evaluation system described in Section 6.The evaluation system comprised 16 DGX-2 nodes and 256 Tesla V100 GPUs.Source: §A.3.2 · p. 13, §A.2 · p. 13

A.3.3 Software Dependencies; A.4 Installation

• Experiments require Ubuntu 20.04, Python 3.8+, and CUDA 11.0+; installation instructions direct users to the artifact READMEs for standalone and integration prerequisites.Integration prerequisites include building PyTorch and installing dependencies for Megatron-LM and NVIDIA Bert.Source: §A.3.3 · p. 13, §A.4 · p. 13

A.3.4 Data Sets.

• Standalone benchmarks need no dataset, whereas the integration experiments require datasets obtained through the coconet-experiments README instructions.The checklist names BERT, GPT-2, and GPT-3 training datasets for integration experiments.Source: §A.3.4 · p. 13, §A.2 · p. 13

A.5.1 Standalone Experiments.

• Standalone experiments generate Figures 10–12 in one hour combined, and MPI’s NPROC setting must not exceed the system’s GPU count.The workflow covers data-, model-, and pipeline-parallel runs and generates their corresponding graphs.Source: §A.5.1 · pp. 13–14

A.5.2 Integration Experiments.

• Integration scripts run data-parallel training, model-parallel inference, and pipeline-parallel inference, taking one hour, under 30 minutes, and three hours, respectively.The data-parallel script reports maximum batch sizes for the 256-GPU evaluation system; the appropriate batch size may differ on another system.Source: §A.5.2 · p. 14

A.6 Evaluation and Expected Results

• Expected standalone outputs match Figures 10–12, while integration outputs match the results in Sections 6.1.2, 6.2.2, and 6.3.2.Source: §A.6 · p. 14