TPU Pod Training

[opooladz]

Hello,

I'm trying to pre-train a llama model using fabric on a TPU Pod. I have access to a few v4-32s. training on a v4-8 is trivial using pytorch XLA on TPUs but scaling to a pod is giving me issues.

LitGPT seems like the most promising pytorch XLA based framework to go with. Can you guys help me with this?

Thanks, Omead