This... doesn't seem to do anything special? Everyone already knew it was bad to "batch" ODEs by making them bigger, e.g. in "Neural Ordinary Differential Equations" (the paper that introduced neural ODEs):
> One can still batch together evaluations through the ODE solver by concatenating the states of each batch element together, creating a combined ODE with dimension D × K. In some cases, controlling error on all batch elements together might require evaluating the combined system K times more often than if each system was solved individually. However, in practice the number of evaluations did not increase substantially when using minibatches.
I don't understand why someone wrote a 30-page, obfuscated paper on just... parallelizing it the obvious way.
I mean, it at least must not be obvious to the poster that says "the claim does sound too good to be true". But yes anyone with enough of an HPC background can look at how vmap is doing its parallelization and instantly know that ML frameworks like Jax and PyTorch are most likely losing an order of magnitude or two of performance. And of course we are very explicit in the paper that this is not novel because we show that the kernels that we are generating match the performance of MPGOS, which is a CUDA library which has the same architecture.
But of course, all of this discussion leaves off half of the title of the paper, "on Multiple GPU Platforms". The point is not that we are able to generate kernels which are doing the fast thing that a dedicated CUDA library does (i.e. not the slow thing that ML libraries are doing), rather the point is that we are doing this in a way where CUDA is not special. We generate similarly optimized kernels for AMD GPUs, Intel GPUs, and Apple silicon (Metal) using this approach. Mose also showed this same codebase can generate kernels for GraphCore IPUs without modifications too (see https://www.youtube.com/watch?v=-fxB0kmcCVE). Showing matching performance with good kernel codes was step 1 but portability (with a full feature set) is step 2. I'd be interested to know if you have any examples of ODE solvers which achieve this level of performance portability because we were unable to find one in the literature or open source.
> One can still batch together evaluations through the ODE solver by concatenating the states of each batch element together, creating a combined ODE with dimension D × K. In some cases, controlling error on all batch elements together might require evaluating the combined system K times more often than if each system was solved individually. However, in practice the number of evaluations did not increase substantially when using minibatches.
I don't understand why someone wrote a 30-page, obfuscated paper on just... parallelizing it the obvious way.