Before anything else, you'll likely want to apply for access to the TPU Research Cloud (TRC). Combined with a Google Cloud free trial, that should allow you to do everything here for free. Once you're in TRC, you need to create a project and with the name of the new project fill out the form that was emailed to you. Use create_tfrecords.py from the GPT-NEO repo to prepare your data as tfrecords; I might do a separate guide on that. Another thing you might want to do is fork the mesh-transformer-jax repo to make it easier to add and modify the config files.
-
Install the Google Cloud SDK. We'll need it later.
-
If you didn't make a project and activate TPU access through TRC yet (or if you plan on paying out of pocket), make one now.
-
TPUs use Google Cloud buckets for storage, go ahead and [cr