Generative AI Google Brain

open-models-serving-vertex-ai-tgi-gemma-multi-lora-adapters-deployment

3 min read

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

Source notebook

Repo path: open-models/serving/vertex_ai_tgi_gemma_multi_lora_adapters_deployment.ipynb · Open on GitHub · advanced

Deploys Gemma 2 with SQL and code LoRA adapters on Vertex AI using a Hugging Face TGI custom handler.

Summary

The notebook teaches how to serve Gemma 2 2B instruction model from Hugging Face Hub with multiple LoRA adapters on Vertex AI. It demonstrates downloading the base model and adapters, testing adapter routing locally with Gemini, writing a custom handler, uploading artifacts to Cloud Storage, registering and deploying the model, and sending online predictions by Python, gcloud, or cURL.

Key code patterns

Initialize Vertex AI

import vertexai
 
vertexai.init(
    project=PROJECT_ID,
    location=LOCATION,
    staging_bucket=BUCKET_URI,
)

Sets the project, region, and staging bucket used for model upload and deployment.

Load base model and adapters

tokenizer = AutoTokenizer.from_pretrained(gemma_path)
model = AutoModelForCausalLM.from_pretrained(
    gemma_path,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16,
    device_map="auto",
)
model.load_adapter(sql_adapter_path, adapter_name="sql_adapter")
model.load_adapter(magicoder_adapter_path, adapter_name="magicoder_adapter")

Loads one Gemma base model and attaches task-specific LoRA adapters for routing.

Route prompts with Gemini

router_model = GenerativeModel("gemini-2.0-flash")
response = router_model.generate_content(
    router_prompt.format(prompt=prompt),
    generation_config=GenerationConfig(
        response_mime_type="application/json",
        response_schema=response_schema,
    ),
).text
classification = json.loads(response)["classification"]

Uses structured JSON output to classify each request as SQL or CODE before selecting an adapter.

Select adapter and generate

if prompt_classification == "SQL":
    model.set_adapter("sql_adapter")
else:
    model.set_adapter("magicoder_adapter")
 
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
with torch.inference_mode():
    generation = model.generate(input_ids=input_ids, generation_config=TGenerationConfig(**generation_config))

Switches adapters per request and runs inference with chat-template tokenization.

Upload and deploy model

model = Model.upload(
    display_name="google--gemma2-tgi-multi-lora-model",
    artifact_uri=str(serve_uri),
    serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-inference-cu121.2-3.transformers.4-46.ubuntu2204.py311",
    serving_container_ports=[8080],
)
deployed_model = model.deploy(
    endpoint=Endpoint.create(display_name="google--gemma2-tgi-multi-lora-endpoint"),
    machine_type="g2-standard-4",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
)

Registers the custom-handler artifact and deploys it on a GPU-backed Vertex AI endpoint.

Models & APIs used

  • Models: google/gemma-2-2b-it, google-cloud-partnership/gemma-2-2b-it-lora-sql, google-cloud-partnership/gemma-2-2b-it-lora-magicoder, gemini-2.0-flash
  • APIs / services: Vertex AI, Cloud Storage, Artifact Registry, Hugging Face Hub
  • SDKs / libraries: google-cloud-aiplatform, vertexai, huggingface_hub, transformers, torch, etils, requests

When to use this

Use this pattern when serving one Hugging Face LLM on Vertex AI with multiple LoRA adapters selected dynamically per prompt.

Gotchas & caveats

  • Gemma weights are gated and require accepting the Google usage license plus a Hugging Face read-only token.
  • The Google Cloud project must enable aiplatform.googleapis.com and artifactregistry.googleapis.com.
  • The service account needs Vertex AI User, Artifact Registry Reader, and Storage Object Admin roles.
  • A valid Cloud Storage bucket must be supplied for intermediate artifacts.
  • Local generation may require about 10 minutes depending on runtime.
  • Vertex AI deployment can take around 15 to 25 minutes.
  • The notebook recommends a GPU runtime and deploys with g2-standard-4 plus one NVIDIA_L4 accelerator.
  • The custom handler string shown references os.getenv and self.tokenizer but imports os is absent and the tokenizer is assigned to self.processor.

Best practices

  • Test the custom handler logic locally before packaging handler.py for serving.
  • Use response_mime_type application/json and a response_schema for the Gemini router output.
  • Validate that each prediction instance contains the inputs key.
  • Use gsutil -m and parallel_composite_upload_threshold=150M for large model artifact uploads.
  • Store handler.py and requirements.txt with the model artifacts in Cloud Storage.
  • Use Vertex AI Model Registry before deploying to an endpoint.
  • Provide Python, gcloud, and cURL prediction paths for endpoint access.
  • Include cleanup flags for endpoint, model, and bucket resources.

Graph View

Backlinks