I am interested in running LLMs like LLAMA-2 and Falcon on Cloud TPUs. These models are trained using PyTorch and not PyTorch XLA or JAX. I am looking for repositories, websites, or tutorials that can guide me in converting these models to run on GPUs instead of TPUs. I found a repository ayaka14732/llama-2-jax which converts LLAMA-2 to JAX, but it’s specific to LLAMA-2. I am seeking more general solutions that can be applied to other LLMs as well. Any suggestions would be greatly appreciated.

I tried running these models on TPU following this guide https://pytorch.org/xla/master/ but even this didn’t work, any help would be greatly appreciated.

Thanks!!