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

arXiv: https://arxiv.org/abs/2105.05720 · HTML: https://mlsanity.com/paper/2105.05720/

## Key facts

- Cited by 40 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 54 references that are in the map.
- First posted on arXiv 12 May 2021.

## Builds on (most cited first, 16 of 16)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Caffe: Convolutional Architecture for Fast Feature Embedding](https://mlsanity.com/paper/1408.5093/) (2014)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [TVM: An Automated End-to-End Optimizing Compiler for Deep Learning](https://mlsanity.com/paper/1802.04799/) (2018)
- [Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms](https://mlsanity.com/paper/2007.00156/) (2020)
- [Loop Tiling in Large-Scale Stencil Codes at Run-time with OPS](https://mlsanity.com/paper/1704.00693/) (2017)

## Built on by (most cited first, 30 of 40)

- [Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve](https://mlsanity.com/paper/2403.02310/) (2024)
- [FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion](https://mlsanity.com/paper/2406.06858/) (2024)
- [Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts](https://mlsanity.com/paper/2502.19811/) (2025)
- [POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference](https://mlsanity.com/paper/2410.18038/) (2024)
- [MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism](https://mlsanity.com/paper/2504.02263/) (2025)
- [MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production](https://mlsanity.com/paper/2505.11432/) (2025)
- [Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference](https://mlsanity.com/paper/2405.17245/) (2024)
- [On Optimizing the Communication of Model Parallelism](https://mlsanity.com/paper/2211.05322/) (2022)
- [FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models](https://mlsanity.com/paper/2501.10714/) (2025)
- [T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives](https://mlsanity.com/paper/2401.16677/) (2024)
- [Optimus-CC: Efficient Large NLP Model Training with 3D Parallelism Aware Communication Compression](https://mlsanity.com/paper/2301.09830/) (2023)
- [TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference](https://mlsanity.com/paper/2505.11329/) (2025)
- [FlashMoE: Fast Distributed MoE in a Single Kernel](https://mlsanity.com/paper/2506.04667/) (2025)
- [Improving Parallel Program Performance with LLM Optimizers via Agent-System Interfaces](https://mlsanity.com/paper/2410.15625/) (2024)
- [MSCCL++: Rethinking GPU Communication Abstractions for AI Inference](https://mlsanity.com/paper/2504.09014/) (2025)
- [Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping](https://mlsanity.com/paper/2409.15241/) (2024)
- [Ladder-residual: parallelism-aware architecture for accelerating large model inference with communication overlapping](https://mlsanity.com/paper/2501.06589/) (2025)
- [A 4D Hybrid Algorithm to Scale Parallel Training to Thousands of GPUs](https://mlsanity.com/paper/2305.13525/) (2023)
- [An Inquiry into Datacenter TCO for LLM Inference with FP8](https://mlsanity.com/paper/2502.01070/) (2025)
- [gLLM: Global Balanced Pipeline Parallelism System for Distributed LLM Serving with Token Throttling](https://mlsanity.com/paper/2504.14775/) (2025)
- [UniEP: Unified Expert-Parallel MoE MegaKernel for LLM Training](https://mlsanity.com/paper/2604.19241/) (2026)
- [MPC-Pipe: an Efficient Pipeline Scheme for Secure Multi-party Machine Learning Inference](https://mlsanity.com/paper/2209.13643/) (2022)
- [TACOS: Topology-Aware Collective Algorithm Synthesizer for Distributed Machine Learning](https://mlsanity.com/paper/2304.05301/) (2023)
- [Tensor-Parallelism with Partially Synchronized Activations](https://mlsanity.com/paper/2506.19645/) (2025)
- [ISO: Overlap of Computation and Communication within Seqenence For LLM Inference](https://mlsanity.com/paper/2409.11155/) (2024)
- [Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment](https://mlsanity.com/paper/2412.07894/) (2024)
- [Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Training](https://mlsanity.com/paper/2601.17654/) (2026)
- [ForestColl: Throughput-Optimal Collective Communications on Heterogeneous Network Fabrics](https://mlsanity.com/paper/2402.06787/) (2024)
- [FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training](https://mlsanity.com/paper/2510.00207/) (2025)
- [CUCo: An Agentic Framework for Compute and Communication Co-design](https://mlsanity.com/paper/2603.02376/) (2026)

## Abstract (arXiv)

Recent trend towards increasing 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, current logical separation between computation and communication kernels in deep learning frameworks misses the optimization opportunities across such barrier. Breaking this abstraction with a holistic consideration can provide many optimizations to provide performance improvements in distributed workloads. Manually applying these optimizations needs modifications in underlying computation and communication libraries for each scenario, which is time consuming and error-prone. Therefore, we present CoCoNeT, with a DSL to express a program with both computation and communication. CoCoNeT contains several machine learning aware transformations to optimize a program and a compiler to generate high performance 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. CoCoNeT enables us to optimize data-, model-and pipeline-parallel workloads in large language models with only a few lines of code. Experiments show CoCoNeT significantly outperforms state-of-the-art distributed machine learning implementations.