Improving Computation–Communication Overlap in Multi-GPU Machine Learning: A Resource-Centric Approach - Using Green Contexts and Tensor Memory Accelerator for Resource-Aware Multi-GPU Communication

Abstract

The rapid growth of deep learning (DL) model size and parameter count has made single-GPU training increasingly impractical due to limited memory and computational capacity. Consequently, large-scale DL models are typically across multiple GPUs. However, multi-GPU execution introduces additional inter-GPU communication overhead, which can become a major performance bottleneck. One common approach to mitigate this overhead is to overlap communication with computation. In such approaches, partial results produced during computation are communicated immediately, while the remaining computation continues in parallel. Since both computation and communication consume GPU resources, efficient overlap requires careful resource allocation. This raises an important question: how should GPU resources be allocated between the computation and communication paths so that overlap can be performed efficiently? This thesis addresses this question by investigating GPU resource utilization during computation-communication overlap, with a particular focus on the GPU’s main execution resources, Streaming Multiprocessors (SMs). Building on prior work, FlashOverlap, the thesis proposes two methods, GCOverlap and Oh Overlap. GCOverlap uses CUDA Green Contexts to partition SMs between computation and communication, reducing interference between the two execution paths. Oh Overlap centers on the communication side by using the Tensor Memory Accelerator (TMA) to reduce the SM resources needed for collective communication. The evaluation shows that GCOverlap can improve operator-level performance for several GEMM+AllReduce and GEMM+ReduceScatter workloads by isolating computation and communication on separate SM partitions. The results also show that Oh Overlap achieves higher bandwidth than its NCCL counterpart for large messages while using fewer SM resources. In overlapping execution, Oh Overlap achieves speedups between 1.1× and 1.4× over the non-overlap baseline, outperforming both the baseline and the original FlashOverlap design in the tested workloads. Overall, this thesis demonstrates that computation–communication overlap is not only a scheduling problem, but also a GPU resource-management problem. Reducing contention and lowering the resource cost of communication can improve overlap efficiency and provide a useful direction for future multi-GPU training systems.