Changes Required for a TPU Pod

[agemagician]

Hello,

Is there any kind of changes that need to be done on the Flax/Jax code to make it run on a TPU pod rather than a single TPU node ?

for example the T5 example on HuggingFace:
https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_t5_mlm_flax.py
Runs normally on a single TPU Pod, but if we need to run it on a TPU Pod, do we need to make any kind of changes to the code?

or we just need to run the same script across all workers and it will communicate with each other:
gcloud alpha compute tpus tpu-vm ssh tpu-vm --zone=us-central2-b --worker=all --command "python myscript.py"

I am trying to run HuggingFace code on TPU V4-64 pod and your reply will be really helpful :)

[agemagician]

The reason for this question is that several users who have access to TPU V4 have a big issue.

The training script for hugging face runs at the same speed and global batch size in a single node or a pod.
For example. running flax t5 in a TPU V4-64 pod reduces the per device batch size to 4 and makes the global batch size to 128, while In a single node it increases to 32 per device and makes the global batch size to 128.
The training speed is exactly the same in a single node or a pod.

[jheek]

The first thing to do here would be to add a log message to your script printing the output of jax.device_count() or even better jax.devices() so you can make sure all devices are discovered by the TPU client. But this does sound like a bug where the different TPU hosts cannot find each other. It's best to report this as an issue in the JAX repo

[dptam]

Hi @agemagician, I was also running into an issue with jax on TPU V4, but when I ran jax.devices, it only showed the CPU. I can run it on a TPU V3 and it will show all the devices. I was wondering if you did anything specific to set it up? Thanks

[agemagician]

Hi @agemagician, I was also running into an issue with jax on TPU V4, but when I ran jax.devices, it only showed the CPU. I can run it on a TPU V3 and it will show all the devices. I was wondering if you did anything specific to set it up? Thanks

I don't have this problem, I just removed jax and installed the latest version and it worked fine.

[agemagician]

It actually discovers all devices and it prints them correctly.
Ok, then I will report it in the Jax repo.
Thanks.

[agemagician]

Yep, that makes a lot of sense.
Thanks a lot :)

[agemagician]

I will ping the HF team to be aware of this issue.

[congyingxia]

Hi @agemagician, do you have example code that can be shared for solving this issue? I met the same problem when I want to adapt the code on TPU-V3-8 to TPU-V3-32. I've read your discussions, but still not sure how to fix it. Do you use the GlobalDeviceArray? Thank you in advance!

[agemagician]

Hi @agemagician, do you have example code that can be shared for solving this issue? I met the same problem when I want to adapt the code on TPU-V3-8 to TPU-V3-32. I've read your discussions, but still not sure how to fix it. Do you use the GlobalDeviceArray? Thank you in advance!

Hi @congyingxia ,
You can check the required change in my pull request which was already integrated into HF for T5 model:
https://github.com/huggingface/transformers/pull/16527/files

[congyingxia]

Thank you so much!

[agemagician]

Thanks again for your effort.