Generative AI Google Brain

open-models-serving-vertex-ai-text-generation-inference-gemma

3 min read

Hugging Face DLCs: Serving Gemma with Text Generation Inference (TGI) on Vertex AI

Source notebook

Repo path: open-models/serving/vertex_ai_text_generation_inference_gemma.ipynb · Open on GitHub · intermediate

Deploys gated Gemma 7B IT from Hugging Face to Vertex AI using a TGI deep learning container.

Summary

This notebook teaches how to register a Hugging Face Hub LLM in Vertex AI Model Registry, deploy it to a GPU-backed Vertex AI endpoint, and send online prediction requests. It configures the Hugging Face TGI container with Gemma model environment variables and a Hugging Face token, then demonstrates prediction via the Vertex AI SDK, gcloud, and cURL. It also covers authentication, IAM/API requirements, chat-template formatting, deployment timing, and cleanup.

Key code patterns

Initialize Vertex AI

from google.cloud import aiplatform
 
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
aiplatform.init(project=PROJECT_ID, location=LOCATION)

Sets project and region before using Vertex AI model registry, endpoints, and prediction.

Register TGI model

model = aiplatform.Model.upload(
    display_name="google--gemma-7b-it",
    serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-2.ubuntu2204.py310",
    serving_container_environment_variables={"MODEL_ID": "google/gemma-7b-it", "HUGGING_FACE_HUB_TOKEN": get_token()},
    serving_container_ports=[8080],
)
model.wait()

Imports the Hugging Face model configuration into Vertex AI using the TGI DLC container.

Deploy endpoint

deployed_model = model.deploy(
    endpoint=aiplatform.Endpoint.create(display_name="google--gemma-7b-it-endpoint"),
    machine_type="g2-standard-4",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
)

Allocates GPU-backed Vertex AI resources for low-latency online prediction.

Apply chat template

tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it", token=get_token())
messages = [{"role": "user", "content": "What's Deep Learning?"}]
inputs = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

Formats instruction-following Gemma prompts before sending them to TGI.

Predict with parameters

output = deployed_model.predict(instances=[{
    "inputs": inputs,
    "parameters": {"max_new_tokens": 20, "do_sample": True, "top_p": 0.95, "temperature": 1.0},
}])

Shows Vertex AI online prediction payload format for TGI text generation.

Models & APIs used

  • Models: google/gemma-7b-it
  • APIs / services: Vertex AI, Artifact Registry
  • SDKs / libraries: google-cloud-aiplatform, huggingface_hub, transformers, jinja2

When to use this

Use this pattern to serve a gated Hugging Face Gemma text-generation model on Vertex AI with the Hugging Face TGI DLC.

Gotchas & caveats

  • Gemma is gated and requires accepting Google’s usage license on Hugging Face Hub.
  • A Hugging Face read-only token is required and passed as HUGGING_FACE_HUB_TOKEN.
  • Vertex AI API and Artifact Registry API must be enabled.
  • Required IAM roles include Artifact Registry Reader and Vertex AI User.
  • Colab may require a runtime restart after package installation.
  • Deployment can take around 15 to 25 minutes.
  • Instruction-following prompts need chat-template formatting because Messages API support is noted as upcoming for TGI 2.3.
  • Machine type, accelerator type, and accelerator count must be compatible.

Best practices

  • Use environment variables for PROJECT_ID and GOOGLE_CLOUD_REGION with us-central1 as the default region.
  • Use a read-only or fine-grained Hugging Face token for google/gemma-7b-it.
  • Set TGI runtime variables such as NUM_SHARD, MAX_INPUT_TOKENS, MAX_TOTAL_TOKENS, and MAX_BATCH_PREFILL_TOKENS.
  • Apply the tokenizer chat template before sending prompts to an instruction-following model.
  • Undeploy endpoints, delete endpoints, and delete registered models to avoid unnecessary costs.

Graph View

Backlinks