Benchmark Tests for FNO and DeepONets

[ayushinav]

To fix #13

• comparison with neuraloperators
• comparison with deepxde
• comparison with Flux variant
• GPU comparisons

I had some issues with CUDA and all when installing torch CUDA toolkit.

[ayushinav]

I guess a more appropriate place for this would be SciMLBenchmarks? @avik-pal

[avik-pal]

Yes the CPU ones can go to SciMLBenchmarks. How are the julia native ones so slow? Did you run the profiler to see where the bottlenecks are?

[ayushinav]

For DeepONet, most of the time goes to the layers, iiuc

For FNOs, ig ffts are a bit expensive

[avik-pal]

Couple of things to check:

1. Load MKL.jl before calling the neural networks
2. Is the number of threads consistent between julia and python?

[avik-pal]

Also that compile function might be compiling the torch model with dynamo. It is pretty much impossible to beat that with Lux running in eager mode. You could try and compile the DeepONet model with EnzymeAD/Reactant.jl#55 and see the performance.

[avik-pal]

Also looking at the plots you are on quite an old version of LuxLib, update it, and it should address some performance

[avik-pal]

@ayushinav can you install LuxDL/LuxLib.jl#111 and let me know how the performance is?

[avik-pal]

I checked the recent lux releases. The current problems are

1. __project isn't fast enough. This needs to be rewritten to use LoopVectorization on CPU
2. batched_mul -- I am looking into this in LuxLib, this is quite easy to fix
3. allocations are expensive. Nothing much we can do there honestly.

The current numbers for Lux on this PR are single threaded, Pytorch uses all cores by default.

[ChrisRackauckas]

Can we make this into a SciMLBenchmarks script? That will be easier to maintain in the long run.

We can make that support GPU

[avik-pal]

Haven't looked into the FNO version much, but that will most likely need LuxDL/LuxLib.jl#118 for performance. To summarize the issue there:

1. gelu and friends are surprisingly expensive operations, and fusing them into the fused_dense operations slows down the entire loop. But that is not true in general. For example, tanh/relu/abs, etc. can and should be fused into the main loop for performance. So we need to extend the current implementation with additional traits to see which activations can be fused for performance.
2. But the main perf gain will come from better gradient ops -- even for gelu. Currently, we fail to use LoopVectorization in that case but if we hardcode some of the common cases the performance will significantly improve.

[avik-pal]

@ayushinav can we get this finished?

[avik-pal]

@ayushinav bump

[ayushinav]

@avik-pal
Yes, I'm on this. I need to write the gradients for the new __projects using LoopVectorization. The projection becomes faster then, but there isn't a significant speed up for the overall network.

[avik-pal]

Do that in a separate PR, let's get the benchmarks aligned first. The pytorch ones use a different size it seems.

[ayushinav]

The sizes are now aligned. The python variant of DeepONet only supports 1 eval point (in the unaligned case), and the Flux variant doesn't support batching. To have the same size for inputs, I made the batch size and the eval points same to compare with both the variants.

The Flux variant of FNO only supported a fixed length of kernels, which is fixed here.

The difference in size now is because python uses (batch_size, N) whereas julia uses (N, batch_size)

[avik-pal]

I am guessing the overhead in FNO is currently from fft?

[ayushinav]

Fair share of fft and matmuladd, though the share of later has increased

[avik-pal]

Can you also profile the backward pass for the FNO? I am surprised it is that bad

[ayushinav]

For mse loss, profiling the backward as gradient(ps -> loss(first(model(x, ps, st)), y), ps)

[avik-pal]

using the permuted formulation it is now all just fft time in forward and backward. It is quite surprising that our FFT is so much slower than pytorch.

@ayushinav might be worth giving https://github.com/Taaitaaiger/RustFFT.jl a shot and checking the performance on CPUs