nanoGPT icon indicating copy to clipboard operation
nanoGPT copied to clipboard

How to train nanoGPT using TPU's?

Open kathir-ks opened this issue 1 year ago • 1 comments

How can I train nanoGPT using TPU's? Can I just modify the DDP targeting TPU VM's or need to make changes to the model to make it XLA compilable?

kathir-ks avatar Feb 05 '24 16:02 kathir-ks

need to use pytorch-xla for that, or reimplement in Jax

VatsaDev avatar Feb 13 '24 20:02 VatsaDev