Documentation request for `scan_kind`

[YigitElma]

Hi,

I am solving an ODE where the vector field needs many complex pytree objects as arguments and also creates some of them for internal use. I have noticed an excessive amount of memory usage on CPU for a single step of Tsit5() (similar behaviour is observed for other adaptive time-stepping solvers too), and the compute time is unreasonably long. However, the same integration on the GPU works fine (at least it runs in a normal amount of time).

In this tutorial, I saw the use of scan_kind="bounded" and gave it a try. It made a huge positive difference in compute time and max memory usage without sacrificing the compile time. Unfortunately, I couldn't find much documentation on different options and their expected behavior. Even realising that this was an option was a bit of luck. For a single step of solver, the overal speedup is like ~30x, from 6 minutes (or more) to just ~10 seconds. And the memory usage goes down to ~100 MBs from 14 GBs.

It would be perfect if this option were documented.

Note: As I mentioned, the vector field definition is pretty convoluted, so I cannot share a minimal example but if you are interested the issue happens for this vector field. The integration is called in this function and an example usage can be found in this test.

Package versions

• jax v0.6.2
• diffrax v0.7.0
• equinox v0.12.2

[patrick-kidger]

Hey there!
Is your use-case around just a forward solve, or are you performing both a forward+backward pass?

Either way, what you're describing is surprising to me. The bounded version is basically a hack to work with forward-mode or higher-order autodiff, since JAX doesn't support jvp-of-custom_vjp. It's intentionally not documented for this reason – I don't regard this as stable and would ideally like to remove it 😁 So the ideal case here would be to figure out why your use-case is exhibiting such poor behaviour with the default choice.

[YigitElma]

Thanks for the answer. FWIW, I also tried the scan_kind="lax" and it has the same issue of very high memory usage and long execution time for my use case (only on CPU...). Btw, these flags didn't have any visible difference on GPU runtime/compile time.

Either way, what you're describing is surprising to me.

I am also shocked by the difference. I suspect this is some XLA compiler difference on CPU/GPU for loops.

Is your use-case around just a forward solve, or are you performing both a forward+backward pass?

For my test case, I just evaluate the function, so only the forward pass. But the same underlying integration routine is also used in an optimization setup where I compute derivatives; in that case, I need the backward pass too (depending on the input-output dimensions). Practically, I never plan to run optimization on a CPU, but still, for pre-/post-processing reasons, it would be nice to be able to run the function on CPU and not kill my laptop (excessive memory usage sometimes causes a shutdown).

My current workaround is to use bounded for just function evaluation purposes on CPU, and use the default for heavier computations on GPU.

[YigitElma]

Just in case someone else reads this, here are docs from equinox.while_loop,

- `kind`: The type of while loop that is lowered to underneath. This may either be
        `"lax"`, `"checkpointed"`, or `"bounded"`.

        If `kind` is `"lax"` then the loop is efficiently forward-mode
        autodifferentiable, but does not support reverse-mode autodifferentiation.

        If `kind` is `"checkpointed"` then the loop is efficiently reverse-mode
        autodifferentiable, but does not support forward-mode autodifferentiation, and
        all `buffers` will be write-only whilst inside the loop.

        If `kind` is `"bounded"` then the loop may be both forward- and reverse-mode
        autodifferentiated, but it requires an `int` value for `max_steps`, and as
        `max_steps` grows then time and memory usage will increase.

[patrick-kidger]

I suspect this is some XLA compiler difference on CPU/GPU for loops.

I agree. I'm noodling this one over and I don't have a great answer for the difference. The main differences between these approaches are in the backward pass. But for just the forward pass, then the default scan_kind and scan_kind="lax" are just a lax.while_loop, whilst scan_kind="bounded" is pretty much just a lax.scan.

My only real suggestions are

• to try wrapping your vector field in a jax.checkpoint in addition to using the default scan_kind, and see what that does. This is what scan_kind="bounded" also has under the hood, though in theory it shouldn't affect the forward pass.
• try different jax+jaxlib versions, to see if this changes things via changes in the XLA compiler.

[YigitElma]

max memory usage without sacrificing the compile time.

Actually, I may have mispoke before. It seems like the compilation takes ages for the CPU version (with default scan_kind) like couple minutes. But for GPU, it is only a couple of seconds. So, when I switch to scan_kind="bounded", the compilation time goes down and it is overall faster. For my test case, compilation went down to 1.7 seconds from 212 seconds.

Thanks for your suggestions! I will try them too, but for now scan_kind looks like the easiest fix. It is kind of tricky to find the actual cause because this only happens on CPU. Do you have any near term plans to remove it :D ?

[YigitElma]

• to try wrapping your vector field in a jax.checkpoint

This didn't make a difference, unfortunately.

• try different jax+jaxlib versions

I tried last couple releases, here are the compile times, (the output line Finished XLA compilation of jit(diffeqsolve) in ... seconds of jax.log_compiles()),

• 0.7.2 -> 3.8 seconds
• 0.6.2 -> 200 seconds
• 0.6.1 -> 45 seconds
• 0.6.0 -> 42 seconds
• 0.5.3 -> 36 seconds
• 0.5.0 -> 30 seconds
• 0.4.38 -> 30 seconds

Apart from 0.7.2, all the dependencies except jax are the same version, for 0.7.2 I had to upgrade equinox to the latest version v0.13.1.

The compile times are all over the place, so, this is definitely an XLA/JAX issue, but scan_kind="bounded" happens to solve it for older JAX versions. For reference, "bounded" compile time is 1.7 seconds which is faster than all above. 0.7.2 is good enough, but we were waiting for people to report possible issues with it before upgrading :)

@patrick-kidger I have one last question about this related to diffrax, as far as I can tell, the loop that makes the most difference is here. Correct me if I am wrong but doesn't most Runge-Kutta methods have a static number of stages? For example, Tsit5 is 7 stages. In this case, isn't jax.lax.scan or jax.lax.fori_loop with static bounds (I guess this reduces to scan) overall better options? Is there a special reason why diffrax uses some form of while_loop?

[patrick-kidger]

It is kind of tricky to find the actual cause because this only happens on CPU. Do you have any near term plans to remove it :D ?

No immediate plans, but it's definitely on my list of warts I'd like to remove Some Day™.

The compile times are all over the place

Wow, absolutely. I can only hypothesise that it's something to do with your choice of vector field, which XLA really does not like put inside a while loop for some reason. (This is doubly mysterious as JAX lowers scans to while loops under the hood.)

If you ever get a MWE then this would be worth raising with the JAX folks. Until then I suppose perhaps we should be happy that at least you have a workaround!

but doesn't most Runge-Kutta methods have a static number of stages?

Not quite. The first step will have one more stage. Later steps get to skip this because of the FSAL property.

(We also need to re-evaluate that first stage if we ever pass a discontinuity in the vector field, which is equivalent to restarting the ODE solve at that location.)

That aside, Runge-Kutta methods involve building up a collection of evaluations in a buffer, with the loop both reading from and writing to that buffer. JAX/XLA is pessimistic about such computational constructs, assuming that each write invalidates the entire buffer. (Not appreciating that we only write to an so-far-unused part of it.) This causes it generate spurious copies of that buffer under autodiff.

[YigitElma]

Thank you very much for the explanation! I think at least the issue related to diffrax has a workaround, and the main bottleneck comes from JAX/XLA. If I can create an MWE, I can share it with them.