SpeedyTransforms: Differentiable transforms via custom Enzyme rule #589

maximilian-gelbrecht · 2024-10-17T16:30:34Z

We want to make the spherical harmonics transforms differentiable.

The complete model is supposed to rely on Enzyme for that, Enzyme can deal with many of our functions already by itself in principle. However, the FFTs are not computed by Julia code, but with FFTW and cuFFT. Transforms are also the most performance critical part of the model, so optimising the gradients of the transforms also makes sense.

Enzyme

The Legendre transformation _legendre! Enzyme can handle by itself
We do Fourier transforms consecutively on each latitude ring in _fourier!, the forward version uses rfft plans, the inverse brfft plans. _fourier! isn't normalised, the normalisation is done in _legendre!.
The FFT needs a rule here, the adjoint of a Fourier transform is the unnormalised inverse Fourier transform (so brfft). For real-valued FFT you have an additional scale in the adjoint reflecting the different coefficients in the formula (there's a two in front of the cos/sine terms)
Rules are tested with the test tools from Enzyme

ChainRules/Zygote/…

For uses outside of our model (e.g. Spherical FNOs) it would be great to also have this differentiable by other ADs. These ADs usually don’t support mutation, so that I would define rules for the non-mutating transform(grid, S). I’ll do that a bit later. This might also need some changes to the code.

This PR

This PR implements the Enzyme Rule for _fourier! and tests it.
As far as I can see it, this rule should also work with the GPU version as soon as we have it as it is
We could pre-compute the scaling. It’s a relatively small Int valued matrix. That would make the adjoint marginally faster.
I also introduce a new extension for FiniteDifferences.jl that makes our data structure compatible with the library
I set up a new environment for the tests.
Currently I just add Enzyme, EnzymeRules and EnzymeTestUtils to the main env to make testing a bit easier for me while developing this PR, that will be changed before merging this
- Do we want to do the differentiability test always in the same files as the other tests, or put them in a separate folder?
- Enzyme (and most of this code) could go in an extension. Is this what we want long-term? Probably yes, but I don’t have a good impression yet how many custom rules we have to define. Hopefully not many though.

Testing

I do a manual test comparing with FiniteDifferences of the full spherical harmonics transforms. Those tests pass
I'd love to test the rules directly with EnzymeTestUtils as well, but there are several problems currently
- There's a problem with the FFT plans having an uninitailzied fields. There's a quite hacky way around this though
- There's a problem with the function having complex valued outputs
- The tests are stuck, I waited 30 min without a result.
- I have to look into this a bit further
FiniteDifferences tests can take a while. Going forward do we want to put all gradient correctness checks in the regular CI? The four quite simple tests that are currently in this PR already take 5 mins.

Complex Numbers

Differentiating complex numbers is a bit of a topic for itself. Enzyme also has a quite long explanation of their approach in their FAQs (https://enzymead.github.io/Enzyme.jl/stable/faq/#Complex-numbers). I think for us, we can treat the real and imaginary part basically as separate numbers, as they would be for real SPH (for which coefficients with m<0 correspond to our imaginary part (sort of)). I am not quite sure yet if this works well out of the box with Enzyme or if we have to be more carefully somewhere. It's something we should think about.

…enzyme

maximilian-gelbrecht · 2024-10-22T12:28:10Z

The forward and backward transforms are differentiable via Enzyme!
The tests comparing the results to finite differences pass on my laptop, hopefully they will in the CI as well

@milankl This might be a good point for you to already have a brief look at it. Especially at what I wrote about how we want to organise this code and testing going forward

maximilian-gelbrecht · 2024-10-22T12:55:10Z

There seem to be some issue running those tests in the CI that I don't have on my computer. All tests are passing on my laptop (except for the EnzymeTestUtils one)

milankl · 2024-10-23T07:16:49Z

With the last commit: You think we can have all Enzyme functionality as an extension? Happy to try, but also loading/precompiling it is not as expensive as CUDA I believe?

maximilian-gelbrecht · 2024-10-23T07:20:52Z

I think so. Because the only thing we'd need Enzyme for directly in the package is custom rules. And I don't think we need that many. If we are lucky, those two might even be the only ones (because they are calling the only non-Julia code).

Enzyme.jl installs the C++ Enzyme library. So I thought, maybe good to move it to an extension when it's not needed for every user.

ext/SpeedyWeatherEnzymeExt.jl

milankl · 2024-10-23T07:25:41Z

ext/SpeedyWeatherEnzymeExt.jl

+
+    # compute the actual vjp
+    dfnorthval = zero(f_north.val)
+    dfsouthval = zero(f_south.val)


Because these allocate: Is that because we can't reuse the forward memory for the gradient?

ext/SpeedyWeatherEnzymeExt.jl

test/runtests.jl

milankl · 2024-10-23T07:33:45Z

test/test_transforms_ad_rules.jl

+
+            # seed
+            dspecs = zero(specs)
+            fill!(dspecs, 1+1im)


What does "seed" mean in this context and why is it 1+i?

We do a reverse mode differentiation here. So you propagate the value of the output in reverse direction to determine the sensitivity of output wrt its inputs. The seed is an example output so to say. If you set it to one, you do actually compute directly the gradient of the function. The reverse mode computes the Jacobian vector product. So with a one seed, you just get the gradient. That's only a very short explanation, should I send you some more recourses about that?

As I wrote above in the big post, differentiating complex numbers is a bit of a special topic in itself. In our case we treat the imaginary and real part kind of like separate numbers, so setting both to one made sense to me intuitively. But differentiating complex numbers is something we should still think about a bit.

Yes please! Just checking, because I think of

julia> z = exp(im*π/4) 0.7071067811865476 + 0.7071067811865475im julia> abs(z) 1.0

more as complex "1" than 1 + 1im because it's actually a unit vector

test/test_transforms_ad_rules.jl

milankl · 2024-10-23T07:38:59Z

test/test_transforms_ad_rules.jl

+
+            fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
+
+            @test isapprox(dspecs, fd_jvp[1])


We could also test that the roundtrip transform is identity no?

What roundtrip exactly? I am not sure I understand. You could differentiate through $\mathcal{S}^{-1}(\mathcal{S}(x))$, but that's just identity, so the derivative wrt the input would just vanish. Or is this exactly what you were thinking of?

Vanish=1 no? It's d(S^-1(S(x)))/dx = dx/dx = 1 I thought? And we could test this starting in either space?

Ah yes sure, you are right.

I added this as tests.

Starting in spectral space it works very well
Starting in grid space, the result is sort of close to 1, but not quite. Currently I am not sure what goes wrong there.

The test starting in grid space is everywhere 0.961122 instead of 1. Looks like a normalisation issue. I'll look into this later

I am not surprised. This depends on what data you start with in grid space, which can hold more information than spectral space, so the first roundtrip is usually lossy. But subsequent ones should be fine. I would start in spectral space, transform to grid space then start the test grid->spectral->grid there

We already do that in the tests. But in principal you are right, there's something wrong with the setup for this test. Because the finite difference differentiation yields the same gradient as Enzyme. So it's not an issue of the custom rule.

maximilian-gelbrecht · 2024-10-23T09:11:29Z

The CI is passing for Julia 1.10.

There are some problems with Enzyme for Julia 1.11, but they are working on it.

milankl · 2024-10-23T11:16:38Z

@vchuravy we (well, Max!) are making some progress to differentiate through SpeedyWeather!

milankl · 2024-10-23T14:31:23Z

@swilliamson7 this might be also relevant for you, let us know if you have any more insights / general wisdom!!

swilliamson7 · 2024-10-23T15:17:12Z

Looking through the thread, but rather exciting for SpeedyWeather.jl! In addition to Valentin, maybe @wsmoses would be interested in seeing this?

maximilian-gelbrecht · 2024-10-24T14:41:47Z

I had a very quick check if some non-sensical differentiation through a timestep of a barotropic model is possible.

I get an error from Enzyme and the compiler from the _divergence! function. That's already a bit what I suspected. After the transform itself, the spatial gradient functions are the other challenging bit. I'll look into differentiating them again, when we also make steps towards GPU-ifying those bits. That should be next on the agenda after the transforms.

milankl and others added 24 commits October 11, 2024 16:11

Do not interweave transform

366ea11

propagate nlayers argument from inputs

42f7e9c

Do not store recursion factors

628f529

SpectralTransform.Grid now keyword argument

c3a340d

introduce scratch memory for FFT output because strides don't match

41eec88

Legendre shortcut generalised

394da96

with fused_oddeven_dot

56a82e5

Fourier/Legendre into functions, smaller scratch output for FFTW

823b66f

first version of a fused matvec

e8c01fd

Batched Legendre transform

41433a3

docstrings added

4c47e15

move even +- odd into kernel

1b80380

ismatching with nlayers check

b057f66

grid to spectral versions v0

873f6d0

ad scale implemented

7dd5811

SpectralTransform constructors overhauled

b5d04c2

Forward _legendre!

9b8363f

Legendre transforms nlayers flexible

57539cf

variables renamed to match harmonics meaning of odd/even

e14aca9

function barrier for fourier_batched, fourier_serial

bd27856

tests adapted

594fc8d

start writing rules

9721f6c

Merge remote-tracking branch 'origin/mk/transform' into mg/transform-…

3269b6e

…enzyme

autodiff computing without error

35055ee

maximilian-gelbrecht added transform ⬅️ ➡️ Our spherical harmonic transform, grid to spectral, spectral to grid differentiability 🤖 Making the model differentiable via AD labels Oct 17, 2024

maximilian-gelbrecht added 4 commits October 17, 2024 18:37

docs and changelog

89d355f

started full transform test (still error)

67be3a3

autodiff on transform! w/o error but no correctness check

5702659

adjust tests (still stuck though)

9148ed7

maximilian-gelbrecht added 4 commits October 22, 2024 09:44

JLArray dependency 0.1 instead of 0.1.4

2cc71d7

transform test sort of working

1783662

inverse transform differentiable

94b760f

Merge branch 'mk/transform' into mg/transform-enzyme

97ddc43

maximilian-gelbrecht changed the base branch from mk/transform to main October 22, 2024 12:25

Merge branch 'main' into mg/transform-enzyme

da88939

maximilian-gelbrecht requested a review from milankl October 22, 2024 12:34

maximilian-gelbrecht added 2 commits October 22, 2024 14:38

add KA to tests dep

a69bfea

test typo in ad tests

b3f7a43

maximilian-gelbrecht added 4 commits October 22, 2024 15:10

add NCDatasets to test env

46eb4c8

update tests for EnzymeTestUtils

d803e59

Add abstractffts to test env

ecf76bf

Enzyme rules as extension

7339f1c

milankl reviewed Oct 23, 2024

View reviewed changes

temporarily deactivate 1.11 CI

acee592

maximilian-gelbrecht added 3 commits October 23, 2024 11:36

add transform derivative identitiy test

f9d444a

deactivate stalling EnzymeTestUtils

660c675

add addtional truncation

ec3af2e

maximilian-gelbrecht added 3 commits October 23, 2024 17:31

update to tests

bf27c17

only one grid for CI

e5a6bd1

additional FD test for identity

fc4a793

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

SpeedyTransforms: Differentiable transforms via custom Enzyme rule #589

SpeedyTransforms: Differentiable transforms via custom Enzyme rule #589

maximilian-gelbrecht commented Oct 17, 2024 •

edited

Loading

maximilian-gelbrecht commented Oct 22, 2024

maximilian-gelbrecht commented Oct 22, 2024 •

edited

Loading

milankl commented Oct 23, 2024

maximilian-gelbrecht commented Oct 23, 2024 •

edited

Loading

milankl Oct 23, 2024

milankl Oct 23, 2024

maximilian-gelbrecht Oct 23, 2024 •

edited

Loading

milankl Oct 23, 2024

milankl Oct 23, 2024

maximilian-gelbrecht Oct 23, 2024

milankl Oct 23, 2024

maximilian-gelbrecht Oct 23, 2024

maximilian-gelbrecht Oct 23, 2024

maximilian-gelbrecht Oct 23, 2024 •

edited

Loading

milankl Oct 23, 2024

maximilian-gelbrecht Oct 24, 2024

maximilian-gelbrecht commented Oct 23, 2024

milankl commented Oct 23, 2024

milankl commented Oct 23, 2024

swilliamson7 commented Oct 23, 2024

maximilian-gelbrecht commented Oct 24, 2024


		fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)

		@test isapprox(dspecs, fd_jvp[1])

SpeedyTransforms: Differentiable transforms via custom Enzyme rule #589

Are you sure you want to change the base?

SpeedyTransforms: Differentiable transforms via custom Enzyme rule #589

Conversation

maximilian-gelbrecht commented Oct 17, 2024 • edited Loading

Enzyme

ChainRules/Zygote/…

This PR

Testing

Complex Numbers

maximilian-gelbrecht commented Oct 22, 2024

maximilian-gelbrecht commented Oct 22, 2024 • edited Loading

milankl commented Oct 23, 2024

maximilian-gelbrecht commented Oct 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

maximilian-gelbrecht Oct 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

maximilian-gelbrecht Oct 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

maximilian-gelbrecht commented Oct 23, 2024

milankl commented Oct 23, 2024

milankl commented Oct 23, 2024

swilliamson7 commented Oct 23, 2024

maximilian-gelbrecht commented Oct 24, 2024

maximilian-gelbrecht commented Oct 17, 2024 •

edited

Loading

maximilian-gelbrecht commented Oct 22, 2024 •

edited

Loading

maximilian-gelbrecht commented Oct 23, 2024 •

edited

Loading

maximilian-gelbrecht Oct 23, 2024 •

edited

Loading

maximilian-gelbrecht Oct 23, 2024 •

edited

Loading