🤖 AI Summary
This work addresses the severe communication bottleneck in distributed Fourier Neural Operators (FNOs) at high resolutions, where spectral layers require frequent all-to-all communication. To mitigate this, the authors propose Distributed Truncated Spectral Transform (DTST), which computes only a small set of critical frequency modes locally on each GPU and aggregates them with minimal communication, achieving equivalence to truncated FFT while drastically reducing communication overhead during both training and inference. DTST uniquely integrates local discrete Fourier transforms with efficient collective communication, unifying spatial data parallelism and spectral weight model parallelism. Experiments demonstrate that, across 4–32 GPUs (up to 8 nodes), DTST accelerates forward propagation by 38–64× and training by 37×, reducing communication time from 97% to under 6% of total execution time, with scalability improving further as resolution increases.
📝 Abstract
Fourier Neural Operators (FNOs) learn solution operators for partial differential equations and offer orders of magnitude speedup over traditional numerical solvers at inference time, which makes them attractive surrogates for high-resolution computational physics. Scaling FNOs to high-resolution spatial grids requires distributing the data across GPUs, but the distributed FFT at the core of each spectral layer requires multiple dense all-to-all collectives that communicate the full spatial tensor, only for most coefficients to be discarded immediately. We introduce the Distributed Truncated Spectral Transform (DTST), which reverses this order. Each GPU computes only a small subset of frequency modes used by the spectral convolution locally via a partial DFT, and two collectives combine the results with a payload that depends only on this mode count, not the spatial resolution. DTST produces spectral coefficients identical to the standard distributed FFT with truncation, while providing both spatial data parallelism and spectral weight model parallelism. We present DRIFT, a GPU implementation of DTST for distributed Fourier Neural Operators, using separable per-dimension basis matrices and efficient GPU-to-GPU communication. On a 3D+time FNO across 4--32 GPUs, on up to 8 nodes (4 GPUs/node), DRIFT achieves a forward-pass speedup of 38--64$\times$ and a 37$\times$ training speedup over the distributed FNO baseline, reducing communication time from 97\% to under 6\% of the forward-pass time, with growing speedups at higher resolution.