search

Lightning-AI / litgpt

20+ high-performance LLMs with recipes to pretrain, finetune and deploy at scale.
https://lightning.ai
Apache License 2.0
9.69k stars 966 forks source link

XLA Pod Request #1100

Open opooladz opened 5 months ago

opooladz commented 5 months ago

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

carmocca commented 5 months ago

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

opooladz commented 5 months ago

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.