Near-Optimal Wafer-Scale Reduce

TL;DR

This paper systematically investigates and optimizes Reduce and AllReduce communication algorithms on the Cerebras Wafer-Scale Engine, achieving near-optimal performance and outperforming existing solutions significantly.

Contribution

It introduces a performance model, new algorithms tailored for WSE, and a code generation approach that enhances Reduce and AllReduce efficiency on wafer-scale architectures.

Findings

Algorithms outperform vendor solutions by up to 3.27x.

Performance model predicts execution time with less than 4% error.

Near-optimal algorithms extend HPC application performance on WSE.

Abstract

Efficient Reduce and AllReduce communication collectives are a critical cornerstone of high-performance computing (HPC) applications. We present the first systematic investigation of Reduce and AllReduce on the Cerebras Wafer-Scale Engine (WSE). This architecture has been shown to achieve unprecedented performance both for machine learning workloads and other computational problems like FFT. We introduce a performance model to estimate the execution time of algorithms on the WSE and validate our predictions experimentally for a wide range of input sizes. In addition to existing implementations, we design and implement several new algorithms specifically tailored to the architecture. Moreover, we establish a lower bound for the runtime of a Reduce operation on the WSE. Based on our model, we automatically generate code that achieves near-optimal performance across the whole range of input sizes. Experiments demonstrate that our new Reduce and AllReduce algorithms outperform the current vendor solution by up to 3.27x. Additionally, our model predicts performance with less than 4% error. The proposed communication collectives increase the range of HPC applications that can benefit from the high throughput of the WSE. Our model-driven methodology demonstrates a disciplined approach that can lead the way to further algorithmic advancements on wafer-scale architectures.