feat(tpu): add release of optimum tpu 0.2.2

README.md

@@ -9,15 +9,14 @@ The [Google-Cloud-Containers](https://github.com/huggingface/Google-Cloud-Contai
 - Training
   - [PyTorch](./containers/pytorch/training/README.md)
     - GPU
-    - TPU (soon)
+    - TPU
 - Inference
   - [PyTorch](./containers/pytorch/inference/README.md)
     - CPU
     - GPU
-    - TPU (soon)
   - [Text Generation Inference](./containers/tgi/README.md)
     - GPU
-    - TPU (soon)
+    - TPU
   - [Text Embeddings Inference](./containers/tei/README.md)
     - CPU
     - GPU
@@ -32,6 +31,12 @@ The [Google-Cloud-Containers](https://github.com/huggingface/Google-Cloud-Contai
 | us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-training-cu121.2-3.transformers.4-42.ubuntu2204.py310  | [huggingface-pytorch-training-gpu.2.3.0.transformers.4.42.3.py310](./containers/pytorch/training/gpu/2.3.0/transformers/4.42.3/py310/Dockerfile)   | PyTorch   | Training  | GPU         |
 | us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-2.transformers.4-44.ubuntu2204.py311 | [huggingface-pytorch-inference-gpu.2.2.2.transformers.4.44.0.py311](./containers/pytorch/inference/gpu/2.2.2/transformers/4.44.0/py311/Dockerfile) | PyTorch   | Inference | GPU         |
 | us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cpu.2-2.transformers.4-44.ubuntu2204.py311   | [huggingface-pytorch-inference-cpu.2.2.2.transformers.4.44.0.py311](./containers/pytorch/inference/cpu/2.2.2/transformers/4.44.0/py311/Dockerfile) | PyTorch   | Inference | CPU         |
+| us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-tpu.0.2.2.py310   | 
+[huggingface-text-generation-inference-tpu.0.2.2.py310](./containers/tgi/tpu/0.2.2/Dockerfile) | 
+| TGI     | Inference | TPU         |
+| us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-training-tpu.2.5.1.transformers.4.46.3.py310   | 
+[huggingface-pytorch-training-tpu.2.5.1.transformers.4.46.3.py310](./containers/tgi/tpu/0.2.2/Dockerfile) | 
+| PyTorch | Training  | TPU         |
 
 > [!NOTE]
 > The listing above only contains the latest version of each of the Hugging Face DLCs, the full listing of the available published containers in Google Cloud can be found either in the [Deep Learning Containers Documentation](https://cloud.google.com/deep-learning-containers/docs/choosing-container#hugging-face), in the [Google Cloud Artifact Registry](https://console.cloud.google.com/artifacts/docker/deeplearning-platform-release/us/gcr.io) or via the `gcloud container images list --repository="us-docker.pkg.dev/deeplearning-platform-release/gcr.io" | grep "huggingface-"` command.
@@ -48,6 +53,8 @@ The [`examples`](./examples) directory contains examples for using the container
 | Vertex AI | [examples/vertex-ai/notebooks/trl-full-sft-fine-tuning-on-vertex-ai](./examples/vertex-ai/notebooks/trl-full-sft-fine-tuning-on-vertex-ai) | Fine-tune Mistral 7B v0.3 with PyTorch Training DLC using SFT on Vertex AI  |
 | GKE       | [examples/gke/trl-full-fine-tuning](./examples/gke/trl-full-fine-tuning)                                                                   | Fine-tune Gemma 2B with PyTorch Training DLC using SFT on GKE               |
 | GKE       | [examples/gke/trl-lora-fine-tuning](./examples/gke/trl-lora-fine-tuning)                                                                   | Fine-tune Mistral 7B v0.3 with PyTorch Training DLC using SFT + LoRA on GKE |
+| TPU       | [gemma-fine-tuning](https://github.com/huggingface/optimum-tpu/blob/main/examples/language-modeling/gemma_tuning.ipynb) 
+| Fine-tune Gemma 2B with PyTorch Training DLC using LoRA |
 
 ### Inference Examples
 

containers/pytorch/training/tpu/2.5.1/transformers/4.46.3/py310/Dockerfile

@@ -34,8 +34,6 @@ RUN pip install git+https://github.com/huggingface/optimum-tpu.git@v${OPTIMUM_TP
 # Add example
 ADD https://raw.githubusercontent.com/huggingface/optimum-tpu/v${OPTIMUM_TPU}/examples/language-modeling/gemma_tuning.ipynb \
   /notebooks/gemma_tuning.ipynb
-ADD https://raw.githubusercontent.com/huggingface/optimum-tpu/refs/heads/llama_tuning_notebook/examples/language-modeling/llama_tuning.ipynb \
-  /notebooks/llama_tuning.ipynb
 
 # Install Google Cloud Dependencies
 RUN pip install --upgrade --no-cache-dir \

containers/tgi/tpu/0.2.2/entrypoint.sh

@@ -1,7 +1,5 @@
 #!/bin/bash
 
-echo "STARTED THE GCP CONTAINER"
-
 # This is required by GKE, see
 # https://cloud.google.com/kubernetes-engine/docs/how-to/tpus#privileged-mode
 ulimit -l 68719476736