The Big Refactor

[ymahlau]

Hey everyone,

first of all thanks to everyone who is contributing and using fdtdx, I am amazed how much traffic there is here for this project (we just passed 300 stars!)

In this discussion section, I want to outline our future plans for fdtdx, an upcoming refactor (which might take a long while) and ask for community feedback. In #273, I described some of the planned features we want to develop. One of them is #104, which aims to increase the simulation speed of fdtdx, both for single-GPU and multi-GPU simulations.

While doing some research and initial tests, it turns out that the only real viable way for implementing this is to do it in PyTorch instead of JAX. While I generally like JAX for its functional programming style, there are strong indications that an implementation in pytorch would be faster, memory-efficient and probably more user friendly (pytorch is generally considered easier to use than JAX).

Some of my the features this would allow to implement:

• fast multi-gpu (single node / multiple nodes). Pytorch has much better multi-gpu support than JAX
• Compression for gradient computation. When doing autodiff, field components need to be saved at the interface between PML and simulation volume. Currently we cannot compress this well, because in JAX all tensors need have a static size. Pytorch does not necessarily have this restriction, which allows for better compression.
• Memory Swap between CUDA and CPU. To enable large-scale simulations on a single GPU, one could exchange data between cpu and gpu to only hold part of the simulation volume at a time in gpu memory. This would only be marginally slower than holding everything in GPU memory (if implemented well) and allow for really large simulations. This type of behavior is generally not possible to implement in JAX.

Since switching to pytorch would require a major rewrite of basically everything in fdtdx, it would also give an opportunity to make the API better. I believe the best way to go would be to make the API similar to the tidy3d API since they took a lot of thought to design their API and it is generally very nice to use.

If you have any thoughts on this let me know! Do you generally prefer using pytorch or JAX? Do you have any suggestions for adapting the API?

[bruxillensis]

Woah, that would be a big change! Fully on board to help with this + would love an API similar to Tidy3D. This might also extend what types of devices can be supported natively (Metal, AMD).

[txdai]

I am also much, much proficient at torch than Jax - and I agree that the in-place modification of torch would truly mean a lot to optimization. I would also be able to contribute to this rewrite.