litgpt icon indicating copy to clipboard operation
litgpt copied to clipboard

Published 20 hours ago • verified publisher icon Lightning-AI

XLA Pod Request

Open opooladz opened this issue 11 months ago • 2 comments

Thank you for the repo.

I am wondering if a recipe for TPU pods can be added. I have access to v4-32 and want to train a LLaMA model from scratch. Wondering if the repo can be extended for this use case.

Thanks

opooladz avatar Mar 13 '24 00:03 opooladz

We have support for a limited set of scripts at https://github.com/Lightning-AI/litgpt/tree/main/xla. Give it a shot, it should work with v4-32. Some info may be outdated

carmocca avatar Mar 14 '24 00:03 carmocca

I mostly see fine-tuning code in those scripts. But there seems to be a general lack of pytorch xla examples which are full from scratch solutions.

opooladz avatar Mar 14 '24 03:03 opooladz