Add train script that runs on the TPU VM

[nostalgebraist]

Adds a script for training on the TPU VM, similar to the existing device_sample.py, etc.

My use case: I'm fine-tuning the model on a single TPU v3-8. I initially got this working with the existing train.py, but

• I had to work around its expectations about file structure e.g. meta.json
• this way has lots of complexity that is needless here: using ray, using an extra VM outside the TPU, the TPUCluster wrapper, etc.

Note: I slapped this together hastily and it's not perfect. I haven't yet verified it works on this branch, just on a personal dev branch with other changes.

Comments

Step

• Compared to train.py, I changed when step is incremented, to try to make the values in the checkpoint file names match the exact number of gradient steps that have been run.
  • I increment the step after compiling model fn
  • Inside the loop, we do (update weights) --> (increment step) --> (save checkpoint if necessary), whereas train.py saves in between the update and the increment
• As a result of this, the val steps are now incremented by one, i.e. we'll do the first val at step 1 (because we took step 1 compiling model_fn)

Data loader

• The changes to the TFRecord reader class cover 2 issues:
  • On my small finetuning dataset, I'm doing multiple epochs. After save/load, the iterator would recurse infinitely because every file was in the used list that was loaded.
    • Fixed with logic to discard used entries from earlier epochs
  • I added reset() method to try to make val data sampling deterministic. (I have a tiny val dataset)

Script

• Does fine-tuning if tune-model-path is provided, otherwise trains from scratch (just for completeness)
• Compared to train.py, adds some more wandb metrics, print statements, etc. You can take or leave these, they're just what I wanted for my project
• save() is mostly a copy-paste of TPUCluster.save which has some useful saving logic. Ideally this logic would be defined only once, but I really wanted to avoid breaking TPUCluster

[kingoflolz]

Looks great, thanks for unwrapping a few layers of this onion and writing some simple code that trains properly haha. Would you be able to write up the config keys required either in the readme or with a commented fine tune config file?

Would also be good if you could write the rough tok/s you get on a v3-8 somewhere for reference

[kingoflolz]

I'm not sure why this is nessasary, can you reset self.used around the except StopIteration: line?

[nostalgebraist]

Oh, yeah, that's a better idea

Changed it to just do self.reset() in that except block

[kingoflolz]

you should probs check that tpu_size <= cores per replica here

[nostalgebraist]

What situation would cause this check to fail?

(Trying to understand what I'm preventing so I can, e.g., write an error message about it)

[kingoflolz]

sometimes JAX doesn't detect the TPUs properly (and detects a single CPU device), or if there is a future model with more than 8 shards

[nostalgebraist]

Doesn't that mean we check that tpu_size >= cores per replica, not <=?

[kingoflolz]

Oh yes, my mistake

[kingoflolz]

Thanks!