Recent trends towards large machine learning models require both training and inference tasks to be distributed.
Considering the huge cost of training these models, it is imperative to unlock optimizations in computation and communication to obtain best performance.
However, the current logical separation between computation and communication kernels in machine learning
frameworks misses optimization opportunities across this barrier.
Breaking this abstraction can provide many optimizations to improve the performance
of distributed workloads.
However, manually applying these optimizations requires modifying the underlying computation and communication libraries for each scenario, which is both time consuming and error-prone. Therefore, we presentCoCoNet, which
contains (i)  a domain specific language to express a distributed machine learning program in the form of computation and communication operations, (ii)  a set of semantics preserving transformations to optimize the program, and (iii)  a compiler to generate jointly optimized communication and computation GPU kernels.
Providing both computation and communication as first class constructs allows users to work on a high-level abstraction and
apply powerful optimizations, such as fusion or overlapping of communication and computation.CoCoNetenabled us to optimize data-, model- and pipeline-parallel workloads in large language models with only a few lines of code.
Our experiments show thatCoCoNetsignificantly outperforms state-of-the-art distributed machine learning implementations.