Memory-Efficient Training with In-Place FFT Implementation

AI summary

Overview

Research area: Memory-efficient deep learning systems / high-performance implementation of Fourier operators (FFT/Inverse FFT) for neural network training.

Technical level: Advanced. The paper assumes familiarity with the Cooley–Tukey FFT decomposition, butterfly operations, twiddle factors, Hermitian/conjugate symmetry of real-valued spectra, circulant weight matrices, and CUDA kernel-level memory management.

Scope: The paper proposes and evaluates rdFFT, a real-domain fully in-place FFT/IFFT operator that keeps input and output in the same real-valued memory buffer, and measures its memory, runtime, and accuracy impact when used inside circulant-structured fine-tuning of LLaMA2-7B and RoBERTa-large.

What This Paper Is About

FFT-based operators are used in neural networks (for example in FourierFT and Block Circulant Adapter fine-tuning) to make weight transformations cheaper, but existing implementations cannot truly reuse the input memory. Standard complex FFT turns N real numbers into N complex numbers, and real FFT (rFFT) turns N real values into N/2+1 complex values stored in the space of N+2 real numbers, so input and output never match in size.

The goal is a real-domain FFT framework (rdFFT) whose forward and backward transforms operate entirely inside the original N-element real-valued buffer, produce the same output as rFFT, allocate no intermediate tensors, and support the bfloat16 data type that existing libraries such as FFTW and cuFFT lack.

Key Contributions

1. rdFFT operator: The first real-valued, fully in-place Fourier transform framework that preserves input-output memory space consistency and eliminates memory space mismatches, producing the same output as rFFT entirely within the original n real-valued input memory space.

2. New memory layout and butterfly scheme: A layout in which each complex coefficient y_k (1 ≤ k < r/2) stores its real part at index k and its imaginary part at the conjugate-symmetric index r-k, plus a novel butterfly execution scheme for IFFT computation, with the special real-valued points y_0 and y_{r/2} each occupying a single real slot.

3. bfloat16 support: Native support for the bfloat16 data type widely used in modern neural networks, which the paper states existing FFT libraries do not provide.

4. Integration and validation: Integration into neural network models via circulant-structured parameter matrices, validated on real models, achieving zero memory allocation for intermediate tensor computations. The implementation is being discussed for upstreaming into PyTorch (github.com/pytorch/pytorch issue 171022).

Main Findings

• Single-layer memory savings: In single-layer training up to the end of the backward pass, rdFFT consistently used the least peak memory. With D=4096, B=1 and block size p=4096, "ours" used 0.09 MB versus 0.25 MB for fft, 0.21 MB for rfft, 20.31 MB for LoRA, and 144.33 MB for full fine-tuning, a reduction of ×1,531.54 relative to full fine-tuning.

• Small-batch advantage, large-batch caveat for baselines: At small batch sizes all circulant variants beat full fine-tuning and LoRA, but as block size decreases and batch size increases, standard FFT-based circulant layers accrue growing overhead from intermediate tensors. At D=4096, B=256, fft with p=128 reached 551.50 MB (×0.30 versus full fine-tuning, i.e., worse) while "ours" with p=1024 used 20.06 MB (×8.19).

• Full-model LLaMA2-7B totals: Base model memory is 12.61 GB for every method. Full fine-tuning totaled 26.90 GB (gradient 6144.00 MB, others 8.28 GB); LoRA r=32 totaled 18.96 GB; fft p=512 20.81 GB; rfft p=512 19.28 GB; "ours" p=512 and p=1024 totaled 17.93 GB and 17.92 GB respectively. The reported "others" figure for ours p=512 and p=1024 is 5.30 GB.

• Full-model RoBERTa-large totals: Base model memory is 1.33 GB. Full fine-tuning totaled 6.15 GB (gradient 192.00 MB, others 4.63 GB); LoRA r=8 6.24 GB; fft p=256 6.72 GB; rfft p=256 6.13 GB; "ours" p=256, p=512 and p=1024 each totaled 5.77 GB with "others" of 4.44 GB.

• Gradient precision asymmetry: In LLaMA2-7B, gradient memory is approximately twice the trainable-parameter memory because the forward pass uses bf16 while gradients must be stored in float32, since backward computation does not support bf16. RoBERTa-large uses full-precision training, so gradient memory matches parameter memory size.

• Operator-level runtime: Measured on an A800 GPU in FP32, averaged over 1000 runs. Forward: at p=512, fft 0.0246 ms, rfft 0.0195 ms, ours 0.0279 ms; at p=4096, fft 0.0252 ms, rfft 0.0199 ms, ours 0.0687 ms. Inverse: at p=512, fft 0.0325 ms, rfft 0.0450 ms, ours 0.0233 ms; at p=4096, fft 0.0327 ms, rfft 0.0470 ms, ours 0.0503 ms. The inverse transform is faster than the forward one because it reuses the butterfly structure in reverse order, reducing dependencies.

• Operator-level accuracy: Absolute error of ours against the torch.fft.fft baseline is 5.99e-07 (p=512), 5.75e-07 (p=1024), 5.84e-07 (p=4096); relative error is 0.0008, 0.0005, and 0.0018. rfft's absolute errors are 1.88e-07, 1.92e-07, 2.55e-07 with relative errors of 0.0001, 0.0001, 0.0012. Errors remain at the level of floating-point numerical noise.

• Model-level throughput and accuracy: On LLaMA-2-7B with GSM8K, token-level throughput in k tokens/sec was 3.29 for full fine-tuning, 3.36 for LoRA, 1.45 for fft, 1.77 for rfft, and 0.92 (p=512), 0.93 (p=1024), 0.93 (p=4096) for ours. MRPC accuracy on RoBERTa-large was 90.9 (full fine-tuning), 90.2 (LoRA), 90.7 (rfft p=512), 90.0 (ours p=512), 89.7 (rfft p=1024), and 90.3 (ours p=1024); several cells are reported as N/A.

• Trade-off stated explicitly: The authors report that their method has lower throughput than fft and rfft but eliminates all intermediate buffer allocations in both forward and backward passes, which they position as the key advantage for large-scale fine-tuning on limited hardware.

Methodology in Plain English

The team started from the Cooley–Tukey algorithm, which computes an FFT by repeatedly applying butterfly operations to ever-smaller sub-transforms. Butterflies are naturally in-place: each output pair overwrites its input pair. The obstacle is that FFT results are complex while neural network tensors are real and stored in fixed shapes.

Their insight is that for real-valued input, every sub-FFT in the recursion also has a real-valued input, so the conjugate symmetry property (y_{r-k} equals the conjugate of y_k) holds at every recursion level. The two special outputs y_0 and y_{r/2} are purely real and need only one slot each, and the remaining complex values can be split so that the real part sits at index k and the imaginary part at the mirrored index r-k. The paper proves (Proposition 1) that in each stage every conjugate-symmetric pair and its butterfly counterparts form a symmetric four-element group around the block center, so the symmetric layout stays valid from stage to stage.

For the inverse transform the situation is harder, because sub-IFFT outputs are not guaranteed to be real or conjugate-structured. The authors exploit linearity: they run the same butterfly graph with reversed data flow, splitting each conjugate pair into symmetric and anti-symmetric components so the recovery uses only real-domain operations and reuses buffer locations.

They then embedded these operators in circulant-structured weight matrices, where a linear layer y = Cx becomes IFFT(FFT(c) ⊙ FFT(x)) and the gradients with respect to x and c are also expressed as IFFT of conjugate-multiplied FFTs. Because the elementwise product of two conjugate-symmetric spectra is itself conjugate-symmetric, the IFFT input in this setting is always symmetric complex-valued, which is what makes the reverse-order in-place scheme applicable.

Evaluation proceeded in two tiers: a controlled single-layer study on an NVIDIA A100 GPU varying input dimension (D=1024 and D=4096) and batch size (B=1, 16, 256) with varying block size p, and full-model training on LLaMA2-7B with GSM8K (per_device_train_batch_size 2, gradient_accumulation_steps 4) and RoBERTa-large with MRPC (batch size 32), using SGD as the optimizer throughout.

Why This Matters

Impact on research: The paper frames operator-level memory optimization as a lossless, complementary strategy to model distillation, quantization, and pruning, which reduce memory by shrinking parameter count or precision rather than by restructuring arithmetic. It also raises the possibility that automatic differentiation frameworks, which typically discourage in-place operations because of gradient-computation difficulties, can accommodate in-place Fourier operators during training.

Real-world applications:

• Parameter-efficient fine-tuning of large language models where GPU memory is the binding constraint, using circulant-structured adapters instead of LoRA.
• Fourier-based fine-tuning methods such as FourierFT and Block Circulant Adapter, which depend on FFT and IFFT inside the training loop.
• Vision and signal-processing models that run partly or entirely in the Fourier domain, including 2D FFT fine-tuning of large language models, spectral convolution, and approximations of global attention.
• Beampattern synthesis and other settings where an IFFT is applied to hidden representations produced by a network.

Industry relevance: The authors note FFTW and cuFFT require pre-allocated buffers of N+2 real numbers that must be handled during model loading, complicating integration, and that those libraries do not support bfloat16. A drop-in, framework-native in-place operator addresses deployment friction directly; the authors state the implementation is being discussed for upstreaming into PyTorch.

Future Directions

• Extending rdFFT to support broader classes of structured transformations beyond circulant matrices.
• Exploring integration with hardware-aware training frameworks for edge deployment, as the conclusion states.
• Adding decoding logic so the implicit real-valued frequency encoding can be converted into an explicitly usable complex spectrum for spectral filtering or custom frequency-domain operations, which the limitations section identifies as currently requiring extra work.
• Addressing the runtime overhead at larger transform sizes, which the authors attribute to CUDA thread-block limitations requiring synchronization both within and across blocks, given that throughput for theirs (0.92–0.93 k tokens/sec) is below fft (1.45) and rfft (1.77).

Target Audience

Systems and efficiency researchers working on memory-constrained training and parameter-efficient fine-tuning; CUDA and numerical-library engineers interested in in-place complex-to-real and real-to-real transforms; and practitioners deploying Fourier-based adapters on hardware where peak GPU memory is the limiting factor. Readers need prior exposure to FFT decomposition and neural network autograd to follow the method and the memory-layout argument.

Authors’ abstract

Fast Fourier Transforms (FFT) are widely used to reduce memory and computational costs in deep learning. However, existing implementations, including standard FFT and real FFT (rFFT), cannot achieve true in-place computation. In particular, rFFT maps an input of size n to a complex output of size n/2+1, causing dimensional mismatch and requiring additional memory allocation. We propose the first real-domain, fully in-place FFT framework (rdFFT) that preserves input-output memory space consistency. By leveraging butterfly operation symmetry and conjugate properties in the frequency domain, we design an implicit complex encoding scheme that eliminates intermediate cache usage entirely. Experiments on multiple natural language understanding tasks demonstrate the method effectiveness in reducing training memory cost, offering a promising direction for frequency-domain lightweight adaptation.