Are there any structural code changes I need to make when running my jax/flax model notebook on a GCP TPU instance vs. a local macOS cpu?

[minimatest]

I ssh'd into a google cloud tpu vm instance. Do I need to edit my jax/flax code in any way to have the tpu perform the computations (e.g. something like pytorch to device("CUDA")). Or would just running the code on the TPU instance suffice? I guess what I am asking is:

Are there any structural code changes I need to make when running my jax/flax model notebook on a GCP TPU instance vs. a local macOS cpu?

[mattjj]

No, you don't need to make any changes for your code to run. You can be explicit about what platforms/devices to use, but by default JAX will use use the GPU or TPU platform if available. (And it always warns you when no GPU or TPU is found, unless you explicitly request the CPU.)

You can see the default device(s) JAX will run on by printing the output of jax.devices(). The default device is always jax.devices()[0]). (By providing an explicit argument to jax.devices, like jax.devices('cpu'), you can also see what devices are available on the non-default platform. Usually the default platform is the interesting one though.)

One thing to keep in mind is that Cloud TPUs have 8 cores, and so jax.devices() will produce a list of length 8. Unless you're using pmap or pjit (or explicitly using jax.device_put to place data/computation, though that's less common and mainly used e.g. to build pipeline computations), JAX will only use the first of the 8 cores. So to make the most use of your Cloud TPU instance, I recommend using pmap or pjit.

WDYT?

[minimatest]

This makes sense. However, the power of the Google Colab and GCP TPU's is that there are many cores to use. If there are 8 cores and by default we are only using one of them, then we aren't making full use. In this regard, I am having great trouble using pmap to distribute training across the cores.

Are there high level abstractions that handle this automatically? Can I place the update params function under a pjit or pmap decoratar and have the code distributed automatically? If not, would frameworks such as Ray help? What do you recommend? Thanks!