Poor performance for GPU with 100.000 parallel integrations, also slow on CPU

[adrianwix]

Hi Marten,

I am currently working on a paper for my bachellor thesis regarding the computation of basin stability for dynamical systems which requires crazy batches of ODEs. In order to find the best solver I am comparing TorchODE, TorchDiffEq, and Diffrax (from JAX).

I am currently having issues getting TorchODE to work for big payloads, for example Diffrax runs in O(1) independently of the batch size which is quite amazing. When running TorchODE for 5k and 10k I get constant time around 30s but when I run it for 100k I get 300s.

I noticed that I was using the wrong controller PIDController instead of IntegralController which is recommended https://torchode.readthedocs.io/en/latest/step-size-controllers/ and improve the time from 1100 seconds to 600 seconds.

Then I noticed that the GPU memory was not clear between benchmarks runs and now it went down to 310 seconds but I have not being able to improve it from there. It could be that I am f...ing it up somewhere since I am quite new to Pytorch.

You can find there results here https://adrianwix.github.io/pyBasin/benchmarks/solvers/.

The code to generate that is https://github.com/adrianwix/pyBasin/blob/main/benchmarks/solver_comparison/benchmark_solver_comparison.py.

Do you think you could take a look at it? I need to submit my paper (and the benchmarks) by the end of February, so I can only change it before that.

Best regards,
Adrian

[martenlienen]

Hey, I noticed that none of your benchmarks use compilation, neither torch.compile nor the equivalent in JAX. This is essential to achieve high performance in cases where the dynamics are simple and quick to evaluate and the majority of time is spent in the solver as in your benchmark. It is also unexpected that torchdiffeq takes the exact same time on both CPU and CUDA.

Could you clarify why torchode is marked as GPU only? It should work just fine on CPU, though it was not optimized for it. Also, I just noticed that the module comment states torchode (dopri5) - CUDA only (no CPU support) but the code only tests with device="cpu"? Or that test is actually skipped is there is none at all that runs torchode?

Also, is there a specific reason why you run with mixed dtype between y and t on CPU but not GPU? torchode will deal fine with it, but if I remember correctly, the larger dtype will propagate to the solution for torchdiffeq, which is not explicitly dtype-stable.

I am curious to see what the results will be with these improvements!

Also, a code improvement: you can read the device and dtype from y with y.device and y.dtype instead of passing y and device separately.

[adrianwix]

Ey, sorry for not responding faster! Let me partially answer some of the questions:

Hey, I noticed that none of your benchmarks use compilation, neither torch.compile nor the equivalent in JAX. This is essential to achieve high performance in cases where the dynamics are simple and quick to evaluate and the majority of time is spent in the solver as in your benchmark.

In JAX when doing batched and using JIT it only saves around 1s after compilation, it looks like the calculation for batched initial conditions is just slow. https://github.com/adrianwix/pybasin/blob/5828aba1079977e143c180a3f1028d60cdeab508/experiments/solver/experiment_jax_jit.py#L100

It is also unexpected that torchdiffeq takes the exact same time on both CPU and CUDA.

Yeah looking at it again it shows that CPU is faster than GPU for smaller N. https://adrianwix.github.io/pybasin/benchmarks/solvers/. I will look into that

Could you clarify why torchode is marked as GPU only?

Regarding marking as GPU I was getting super high execution times on CPUs so disable it because it looks like there may be a problem. I am trying to remake a CPU benchmark. Here is a snapshot of the CPU run for torchode and torchdiffeq, it is already interesting to see that torchdiffeq works better with nn.Module and torchode with a normal function.

I will send you the file with the benchmark and the results when it is done since I am working on other changes in parallel.

Also, is there a specific reason why you run with mixed dtype between y and t on CPU but not GPU? torchode will deal fine with it, but if I remember correctly, the larger dtype will propagate to the solution for torchdiffeq, which is not explicitly dtype-stable.

Mmm there I am still figuring out things, when trying to achieve better performances I tested with float32 because it should be faster for GPUs but then I think with float32 tight tolerances do not work. I am still figuring out the whole thing. Also checking Julia solver they allow for different typing for time and initial conditions

I will keep looking into this and let you know when I have better news. Thanks for your support!

[martenlienen]

Mmm there I am still figuring out things, when trying to achieve better performances I tested with float32 because it should be faster for GPUs but then I think with float32 tight tolerances do not work. I am still figuring out the whole thing. Also checking Julia solver they allow for different typing for time and initial conditions

torchode also handles it correctly, but torchode does not, if I remember correctly and your state will become float64 if t is float64, even if your y0 is float32.

[adrianwix]

torchode also handles it correctly, but torchode does not

Wait what? You mention torchode twice 😅

[martenlienen]

Sorry, I meant torchdiffeq does not :D

[JiwaniZakir]

The superlinear scaling you're seeing (5k→10k stays flat at ~30s, but 100k jumps to ~300s) is characteristic of TorchODE's shared step size mechanism: all parallel integrations are forced to march at the same step size, determined by whichever trajectory in the batch has the most restrictive error tolerance requirement. With 100k integrations, the probability of having outlier trajectories with stiff or chaotic local behavior is dramatically higher, which can collapse the global step size and force orders of magnitude more steps across the entire batch. This is fundamentally different from how Diffrax achieves O(1) scaling — JAX's vmap-based approach allows each integration to take independent step sizes. The relevant logic in TorchODE is in the step size controller's merge_error_norms or equivalent reduction step, where the minimum acceptable step across the batch is enforced; you can verify this is the bottleneck by logging the number of steps taken per solve and comparing across batch sizes.

[martenlienen]

Thanks for weighing in, but torchode uses the same approach as diffrax, albeit not using vmap. In torchode, each ODE has a separate step size. Also, torchode does not have a function called merge_error_norms. The constant scaling in adrianwix's benchmark is not due to the handling of time steps, but due to CUDA. You can see that torchdiffeq, which handles time steps like you describe above, also has constant time on CUDA.