Ways to load GPT-NeoX checkpoints in GPT-Neo for TPU training?

[frankxu2004]

Thanks for the great project. I wonder if I could load the checkpoints for GPT-NeoX to GPT-Neo for the reason of leveraging TPU training. Or better yet does GPT-NeoX supports TPU training as well?

[StellaAthena]

There is not currently a supported way to convert GPT-NeoX checkpoints to GPT-Neo, and you cannot use GPT-NeoX on a large TPU cluster. I do not anticipate us adding this as a feature in the future either.

For training on TPUs I recommend checking out Mesh Transformer JAX, the library that trained GPT-J.

[nikhilanayak]

Can GPT NeoX 20B be trained with Mesh Transformer JAX?