In order to satisfy their ever increasing capacity and compute requirements, many machine learning models are distributed across multiple nodes using space-efficient parallelism strategies. As a result, collective communications are often on the critical path, and hiding their latency by overlapping kernel-granular communication and computation is difficult due to the absence of independent computation. In this work, we propose fusing computation with communication using GPU-initiated networking, and leverage GPUs' massive parallelism to enable fine-grained overlap of the fused operations. We have developed a single, self-contained GPU kernel where workgroups (WGs) immediately communicate their results to remote GPUs when they complete their computation. Meanwhile, other WGs within the same kernel perform overlapping computation, maintaining high ALU utilization. Furthermore, we propose zero-copy optimizations for peer-to-peer GPU communication where the data computed by one GPU is directly written to the destination buffers within the peer GPUs, eliminating intermediate stores and extra buffering. Our approach leverages the emerging multi-node GPU system trend where GPUs are physically close to network with direct GPU-NIC interconnects. We demonstrate our approach by creating an embedding + All-to-All fused kernel which overlaps embedding operations and the dependent all-to-all collective in DLRM models. We evaluate our approach both using simulation and real hardware. Our evaluations show that our approach can effectively overlap All-to-All communication with embedding computations, subsequently reducing their combined execution time by 31% on average (up to 58%) for inter-node and by 25% (up to 35%) for intra-node configurations. Scale-out simulations indicate that our approach reduces DLRM execution time by ~10% for 128 node system.

翻译：为满足不断增长的容量和计算需求，许多机器学习模型采用空间高效的并行策略跨多个节点分布。因此，集体通信常处于关键路径上，而由于缺乏独立计算，通过内核粒度的通信与计算重叠来隐藏其延迟十分困难。本研究提出利用GPU启动的网络技术将计算与通信融合，并借助GPU的大规模并行性实现融合操作的细粒度重叠。我们开发了一个独立的自包含GPU内核，其中工作组（WGs）在完成计算后立即将结果传输至远程GPU，同时同一内核内的其他工作组执行重叠计算，保持高ALU利用率。此外，我们针对点对点GPU通信提出零拷贝优化，使一个GPU计算的数据可直接写入对等GPU的目标缓冲区，从而消除中间存储和额外缓冲。该方法利用了新兴的多节点GPU系统趋势——GPU与网络物理接近且具备直接GPU-NIC互连。我们通过创建嵌入层与全对全（All-to-All）融合内核来演示该方法，该内核在DLRM模型中重叠嵌入操作及其依赖的全对全集体通信。基于仿真与真实硬件的评估显示，该方法能有效重叠全对全通信与嵌入计算：节点间配置下平均缩短31%（最高58%）的合并执行时间，节点内配置下平均缩短25%（最高35%）。规模扩展仿真表明，在128节点系统中，该方法可将DLRM执行时间减少约10%。