Descriptor-Resident Distributed Spectral Operators for High-Resolution Fourier-Domain Neural Workloads
A descriptor-resident multi-GPU spectral backend for PyTorch focused on correctness, autograd support, strict backend execution, and high-resolution Fourier-domain neural workloads.
Problem
Fourier-domain neural workloads can exceed the memory of a single accelerator while still requiring mathematically correct transforms and differentiable execution.
Motivation
The work examines how a distributed spectral backend can preserve correctness and framework integration when problem sizes move beyond a single GPU.
Architecture
The backend combines Python and PyTorch integration with C++/CUDA execution, cuFFT/cuFFTXt primitives, and distributed coordination through MPI and NVSHMEM.
Descriptor-resident execution
Execution descriptors remain resident so the backend can enforce its intended distributed path rather than silently falling back to an unintended implementation.
Correctness methodology
Validation covers operator parity for FFT2, FFT3, and DCT2, along with automatic-differentiation and backend-execution tests.
Capacity experiments
Capacity observations are scoped to the recorded four-NVIDIA-V100 environment and are not generalized as universal performance or scaling claims.
Experimental environment
Recorded capacity experiments used four NVIDIA V100 GPUs. The reported measurements are environment-specific observations.
Current research status
Research in progress; the implementation and test methodology continue to be developed and evaluated.
Limitations and next questions
This is active research. Broader hardware validation, additional workload characterization, and continued performance analysis remain open research questions.
Personal ownership
Yash is the researcher and primary systems implementer, responsible for distributed systems development and correctness, autograd, backend-execution, capacity, and performance tests.
Recorded experimental evidence
Operator parity
~5e-7 maximum relative error
FFT2, FFT3, and DCT2 parity evaluation in the specific NVIDIA V100 experimental environment.
Verification: Recorded operator parity tests.
FNO-forward capacity
24,576 × 24,576
Tested on four NVIDIA V100 GPUs in the specific V100 experimental environment.
Verification: Recorded capacity experiment.
In-place distributed FFT
64,512 × 64,512
Tested on four NVIDIA V100 GPUs in the specific V100 experimental environment.
Verification: Recorded capacity experiment.