MKernel speeds multi-GPU and multi-node training with fused kernels
mKernel: Fast Multi-GPU, Multi-Node Fused Kernels
Distributed, Parallel, and Cluster ComputingArtificial IntelligenceMachine LearningNetworking and Internet Architecture
Summary
Training large AI models across many graphics processors (GPUs) and computers is slow because of the time spent moving data between them. The authors present mKernel, a software tool that runs calculations and data transfers together efficiently by splitting up GPU resources and carefully managing communication inside and outside the computer. This approach reduces the waiting time for data transfers during training and inference. Tested on powerful GPU clusters, mKernel achieves noticeable speed improvements for common AI computation tasks.
What this means in practice
- •For machine learning engineers: Run faster training and inference of large AI models on multiple GPUs and nodes with reduced communication delays using mKernel fused kernels.
- •For high performance computing teams: Integrate mKernel to improve efficiency of data movement and computation overlap in multi-node GPU clusters with RDMA networks.
Authors
Ziming Mao, Yihan Zhang, Shawn Wei Chew, Shuang Ma, Costin Raiciu, Yang Zhou, Scott Shenker, Ion Stoica
Abstract
Communication has become a bottleneck in distributed training and inference of large models. Overlapping communication with computation at the granularity of kernels, on separate streams, reduces only part of this communication cost. Fused kernels often have better performance by transmitting each output tile as soon as it is produced, but existing fused kernels are largely confined to a single NVLink domain. We present mKernel, a library of multi-GPU, multi-node fused kernels that overlap computation, intra-node NVLink communication, and inter-node RDMA at tile granularity. mKernel partitions the streaming multiprocessors (SMs) of a persistent kernel into compute and communication roles, and an on-GPU controller tunes the SM partition adaptively at run time, since the best SM partition varies with the kernel and the input shape. It structures data movement hierarchically so that data traversing the inter-node network is minimized. Finally, it drives the network from the GPU through a lightweight command queue and host proxy implemented directly on RDMA verbs, which allows the same kernels to run on any network backend (e.g. InfiniBand and on AWS EFA); we observe, surprisingly, that GPUDirect Async (IBGDA) yields little additional benefit over host-assisted GPU-initiated communication. We implement five kernels spanning tensor, sequence, and expert parallelism. On two 16-GPU H200 clusters, mKernel achieves speedups of up to 1.72x on GEMM+AllReduce and $1.88\times$ on Ring Attention.