Generative AI Google Brain

open-models-serving-get-started-with-vllm-lora-serving-on-vertex-ai

3 min read

Deploying Multiple LoRA Adapters on Vertex AI with vLLM

Source notebook

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

Deploys Gemma 2 with multiple LoRA adapters on Vertex AI using a custom vLLM container.

Summary

This notebook teaches how to serve one Gemma 2 base model with multiple preloaded LoRA adapters on a Vertex AI endpoint. It downloads a Hugging Face base model and SQL/code adapters, uploads them to Cloud Storage, builds a vLLM container with a GCS-loading entrypoint, registers and deploys it on Vertex AI, then tests base, SQL, and code generation requests.

Key code patterns

Initialize Vertex AI and bucket

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
BUCKET_URI = f"gs://{PROJECT_ID}-vllm-peft-serving"
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the project, region, storage location, and Vertex AI SDK context used by the deployment.

Download gated Hugging Face assets

interpreter_login()
base_model_path = snapshot_download(repo_id=base_model_id, token=get_token())
sql_adapter_path = snapshot_download(repo_id=sql_adapter_id, token=get_token())
magicoder_adapter_path = snapshot_download(repo_id=magicoder_adapter_id, token=get_token())

Downloads the base Gemma model and LoRA adapters after Hugging Face authentication.

Use GCS paths in vLLM args

vllm_args = [
  "python3", "-m", "vllm.entrypoints.openai.api_server",
  f"--model={BUCKET_URI}/models/gemma-2-2b-it",
  "--enable-lora", "--max-loras=4", "--max-lora-rank=64",
  "--lora-modules",
  f"sql-lora={BUCKET_URI}/adapters/sql-lora",
  f"magicoder-lora={BUCKET_URI}/adapters/magicoder-lora",
]

Preloads multiple LoRA adapters so clients can select an adapter per request.

Upload custom serving container

vertexai_model = aiplatform.Model.upload(
  display_name=MODEL_NAME,
  serving_container_image_uri=DOCKER_URI,
  serving_container_args=vllm_args,
  serving_container_predict_route="/v1/completions",
  serving_container_health_route="/health",
)

Registers a Vertex AI Model resource backed by the custom vLLM container.

Select adapter per request

request_body = json.dumps({
  "model": "sql-lora",
  "prompt": prompt,
  "max_tokens": 150,
  "temperature": 0.0,
})
response = vertexai_endpoint.raw_predict(body=request_body, headers={"Content-Type": "application/json"})

Uses OpenAI-compatible request fields through Vertex AI raw prediction to activate a specific adapter.

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
  • APIs / services: Vertex AI, Cloud Storage, Artifact Registry, Cloud Build
  • SDKs / libraries: google-cloud-aiplatform, vertexai, huggingface_hub

When to use this

Use this pattern when serving several task-specialized LoRA adapters from one open base model on a managed Vertex AI endpoint.

Gotchas & caveats

  • Vertex AI API must be enabled and IAM roles include aiplatform.user, artifactregistry.admin, cloudbuild.builds.editor, and storage.admin.
  • Gemma and adapters may be gated on Hugging Face, requiring account access, accepted terms, and a read token.
  • The custom container downloads models from GCS at startup, so deployment waits for model download and health checks.
  • LoRA support requires —enable-lora and adapter ranks must not exceed —max-lora-rank.
  • Increasing —max-model-len or —max-loras can raise GPU memory pressure and may require lowering gpu-memory-utilization.
  • Streaming through Vertex AI raw_predict needs additional server-side configuration; production streaming may require direct vLLM OpenAI-compatible endpoints.
  • Dynamic runtime adapter loading requires VLLM_ALLOW_RUNTIME_LORA_UPDATING=True, redeployment, and direct vLLM API access, not Vertex AI raw_predict.
  • Cleanup is needed for endpoint, model, Artifact Registry repository, and Cloud Storage bucket to avoid charges.

Best practices

  • Store model files in Cloud Storage instead of baking them into a large Docker image.
  • Enable parallel composite uploads for faster Cloud Storage transfers.
  • Use a startup entrypoint that fails fast when GCS downloads fail before starting vLLM.
  • Set temperature to 0.0 for deterministic SQL generation.
  • Configure Vertex AI health and predict routes explicitly for the serving container.
  • Use autoscaling target accelerator duty cycle for GPU-based endpoint scaling.
  • Delete deployed resources after testing to avoid extra charges.

Graph View

Backlinks