Breaking the computation and communication abstraction barrier in distributed machine learning workloads
Breaking the computation and communication abstraction barrier in distributed machine learning workloads
复制标题
DOI：
10.1145/3503222.3507778
复制
发表时间：
2021-05
期刊：
影响因子：
--
通讯作者：
Abhinav Jangda;Jun Huang;Guodong Liu;Amir Hossein Nodehi Sabet;Saeed Maleki;Youshan Miao;M. Musuvathi-M.-Mu
中科院分区：
文献类型：
--
作者：
Abhinav Jangda;Jun Huang;Guodong Liu;Amir Hossein Nodehi Sabet;Saeed Maleki;Youshan Miao;M. Musuvathi-M.-Mu
最近大型机器学习模型的趋势要求训练和推理任务都是分布式的。考虑到训练这些模型的巨大成本，解锁计算和通信方面的优化以获得最佳性能势在必行。然而，目前机器学习框架中计算和通信内核之间的逻辑分离错过了跨越这一障碍的优化机会。打破这种抽象可以提供许多优化，以提高分布式工作负载的性能。然而，手动应用这些优化需要修改每个场景的基础计算和通信库，这既耗时又容易出错。因此，我们提出了CoCoNet，它包含(I)一种特定于领域的语言来以计算和通信操作的形式表达分布式机器学习程序，(Ii)一组保持语义的转换来优化程序，以及(Iii)一个编译器来生成联合优化的通信和计算GPU内核。将计算和通信都作为第一类构造提供，允许用户进行高级抽象并应用强大的优化，例如通信和计算的融合或重叠。CoCoNet使我们能够在大型语言模型中优化数据、模型和管道并行工作负载，只需几行代码。我们的实验表明，CoCoNet的性能明显优于最先进的分布式机器学习实现。
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 present CoCoNet, 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. CoCoNet enabled us to optimize data-, model- and pipeline-parallel workloads in large language models with only a few lines of code. Our experiments show that CoCoNet significantly outperforms state-of-the-art distributed machine learning implementations.