[TPU] XLA changes for finetuning

[gkroiz]

This PR adds changes specific for finetuning on TPUs.
Used TPU v4-8 with

log_interval = 1
devices = 4
batch_size = 64 / devices
micro_batch_size = 4
gradient_accumulation_steps = 4
num_epochs = 5

for stablelm-base-alpha-3b; resulting finetuning time ~= 4500s.

[gkroiz]

Would it be simpler to have a separate file for tpu finetuning (for example, adapter_tpu.py)?

[gkroiz]

cc @Liyang90 for review

[gkroiz]

TODO: if we decide to keep these changes in adapter.py, we also need to make the same adjustments in adapter_v2.py

[carmocca]

Some early questions!

[gkroiz]

When using TPUs, the current code in adapter.py won't work when using more than 8 cores (4 chips).

For example, let's say we want to run finetune on v4-64 (64 cores, 32 chips). We would need two device counts, one for local devices within each of the 8 workers and one for the total device count. The reasoning here is that the XLA fabric strategy needs to define devices=4 (local device count) but for batch_size, max_iters, and warmup_iters, we would want to define based on the total device count, which would be 32. For my local tests, I made this distinction by defining local_devices and total_devices, but I'm not sure if this leads to unnecessary confusion when using non-xla strategies. What are your thoughts @carmocca?

[carmocca]

I think making these distinctions is fine. Because they are still true in the general multi-node case, not just with XLA.

[gkroiz]

@carmocca I noticed this PR is now closed. Do we no longer need to xla changes?

[carmocca]

Sorry, it was an accident!

[carmocca]

Minor comment. Rest looks good