Multi-GPU communication improved with clear memory rules in device APIs
A New Lens on Multi-GPU Communication: API-Level Memory Models for Portable Optimization
Distributed, Parallel, and Cluster Computing
Summary
Slow communication between GPUs can limit the speed of machine learning and scientific computing. The authors propose clear rules (memory models) for how communication commands work at the API level, so programmers can write faster and more reliable code. They show this approach in a popular GPU library and achieve big speedups in several complex workloads. This helps avoid bugs and compatibility issues when optimizing data movement across GPUs.
What this means in practice
- •For high performance computing teams: Improve data transfer speeds in HPC applications by updating GPU communication APIs with defined memory semantics enabling safer low-level optimizations.
- •For machine learning engineers: Boost training efficiency of large models by using standardized API memory models that allow cache bypassing and faster device communication.
Authors
Johnathan Alsop, Bradford Beckmann, Brandon Potter, Jiakun Yan, Aurelien Bouteiller, Omri Mor, Tony Tye, Sameer Sahasrabuddhe
Abstract
As network communication increasingly becomes a performance bottleneck for ML and HPC workloads, it is more important than ever to optimize multi-GPU data transfers. Recently, low-level device APIs have emerged that allow programmers to perform fine-grain optimizations via one-sided put/get operations triggered directly from the device. However, the communication libraries providing these device APIs (e.g., rocSHMEM, MSCCL++, and NCCL/RCCL GIN) currently lack a rigorous memory model defining how their primitives interact with the underlying memory systems. Without this, emerging HPC and ML workloads aggressively optimizing network communication via device APIs can be brittle, error-prone, and face compatibility and inter-operability challenges. This paper proposes the use of API-level memory models for these emerging device communication APIs. An API-level memory model defines the memory semantics of communication API calls in a way that abstracts lower-level details from the API user and simplifies development. This work defines one such model for device-side communication (implemented in AMD's recently open sourced rocSHMEM library) and demonstrates how formal relaxations to such a model can enable low-level LLM performance optimizations while retaining a well-defined standardized interface. Specifically, we describe how updating the SHMEM API implemented by rocSHMEM can properly support cache bypassing for network data, improving performance of a collective communication workload by up to 2.48x, a fused compute+communication workload by up to 1.77x, and a state-of-the-art expert-parallel communication workload by up to 1.18x.