Reverting the TGI image version for LLAMA multiple GPUs in GKE samples

[raushan2016]

Description

The current image override the HF_HOME to /tmp from /data. Even after changing the mountpath to /tmp there is some regression in the newer TGI image which results into out of GPU memory on L4 and requires atleast A2 node. Rolling back the image version to get the sample working will investigation happen in the background.

Issue: #1581

Tasks

• The contributing guide has been read and followed.
• The samples added / modified have been fully tested.
• Workflow files have been added / modified, if applicable.
• Region tags have been properly added, if new samples.
• Editable variables have been used, where applicable.
• All dependencies are set to up-to-date versions, as applicable.
• Merge this pull-request for me once it is approved.

[code-review-assist]

Hi team,

I'm currently reviewing this pull request and will provide my feedback in a few minutes. In the meantime, here's a summary of the changes for you and other reviewers to quickly get up to speed:

This pull request, authored by @raushan2016, reverts the Text Generation Inference (TGI) image version used in several Google Kubernetes Engine (GKE) samples for large language models (LLMs) running on multiple GPUs. The intent is to resolve a regression introduced by a newer TGI image version that caused out-of-GPU memory errors on L4 nodes, requiring the use of more powerful A2 nodes.

Here's a breakdown of the changes:

• Image Version Reversion: The PR changes the image specification in several YAML files related to LLM samples. The old image versions (us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-3.ubuntu2204.py311) are replaced with older, more compatible versions:
  • ai-ml/llm-multiple-gpus/falcon-40b/text-generation-inference.yaml: Changed from us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-3.ubuntu2204.py311 to ghcr.io/huggingface/text-generation-inference:1.4.3 (lines 32)
  • ai-ml/llm-multiple-gpus/llama2-70b/text-generation-inference.yaml: Changed from us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-3.ubuntu2204.py311 to ghcr.io/huggingface/text-generation-inference:1.4.3 (lines 32)
  • ai-ml/llm-multiple-gpus/llama3-70b/text-generation-inference.yaml: Changed from us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-3.ubuntu2204.py311 to ghcr.io/huggingface/text-generation-inference:2.0.4 (lines 32)
  • ai-ml/llm-multiple-gpus/mixtral-8x7b/text-generation-inference.yaml: Changed from us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-3.ubuntu2204.py311 to ghcr.io/huggingface/text-generation-inference:1.4.3 (lines 32)
• Clarifying Comments: Comments were added to the YAML files to explain the /data mount path's relation to the HF_HOME environment variable (lines 54-56, 59-61, 61-63, 60-62). This improves readability and understanding.

The PR addresses issue #1581. A deeper investigation into the regression in the newer TGI images will happen separately.

---

A quick haiku to lighten the mood:

Code flows like a stream,
Bugs fixed, a smoother journey,
Joy in the review.

[code-review-assist]

This pull request addresses the OOM issues on L4 and A2 nodes by reverting the TGI image versions for falcon-40b, llama2-70b, and mixtral-8x7b to 1.4.3. The llama3-70b image is reverted to 2.0.4 for consistency. The approach seems appropriate given the circumstances, and the added comments explaining the mountPath are helpful. It's crucial to continue investigating the root cause of the OOM issue in the newer TGI images.

[moficodes]

@kenthua I am ok with this change.

[moficodes]

/lgtm

[kenthua]

In syncing with @raushan2016 , we have deeplearning container image that does work for this example.

[kenthua]

can we add the 20241210-1804-rc0 image tag?

[alvarobartt]

Thanks for the PR, added some nits! Also I'll check the issues with both Falcon 40B and Llama 2 70B and come back with another PR on top once those are solved, thanks for your time again! 🤗