search

huggingface / Google-Cloud-Containers

Including Hugging Face Deep learning Containers for Google Cloud
Apache License 2.0
112 stars 10 forks source link

Feature/pytorch tpu container #14

Closed shub-kris closed 5 months ago

shub-kris commented 7 months ago

This PR adds PyTorch TPU Dockerfile.

shub-kris commented 7 months ago

It is based on the image mentioned in the PyTorch_XLA GitHub repo. So, far I have tested it on the example mentioned in Google Cloud TPU docs.

shub-kris commented 7 months ago

Added notebook, as it was not installed, also removed libraries that doesn't support TPU.

shub-kris commented 6 months ago

@philschmid let's merge this PR. I will test with the new transformers version 4.38.1 in a different branch and do a PR then.

shub-kris commented 6 months ago

@philschmid @tengomucho according the discussions in the slack, I have updated the Dockerfile to have the nightly version, also I have ran the finetune-gemma-lora-dolly.py and it runs successfully and saves the checkpoints too.

philschmid commented 6 months ago

Did we validate if the results are good?

shub-kris commented 6 months ago

@tengomucho tried to improve the commit message