Add multi-GPU support

[MUCDK]

As ott-jax supports this now, we can also have it in moscot.