Generative AI Google Brain

open-models-fine-tuning-vertex-ai-trl-fine-tuning-gemma

3 min read

Hugging Face DLCs: Fine-tuning Gemma with Transformer Reinforcement Learning (TRL) on Vertex AI

Source notebook

Repo path: open-models/fine-tuning/vertex_ai_trl_fine_tuning_gemma.ipynb · Open on GitHub · intermediate

Fine-tunes google/gemma-2b with TRL SFT and LoRA in a Vertex AI custom container job.

Summary

This notebook teaches how to run supervised fine-tuning for Gemma using Hugging Face TRL, PEFT LoRA, and a PyTorch training DLC on Vertex AI. The workflow authenticates Google Cloud and Hugging Face, configures project, region, and a GCS bucket, initializes the Vertex AI SDK, creates a CustomContainerTrainingJob, and submits a single-GPU L4 training job that writes artifacts to Cloud Storage.

Key code patterns

Initialize Vertex AI

aiplatform.init(
    project=os.getenv("PROJECT_ID"),
    location=os.getenv("LOCATION"),
    staging_bucket=os.getenv("BUCKET_URI"),
)

Binds the custom training job to a Google Cloud project, region, and staging bucket.

Create TRL container job

job = aiplatform.CustomContainerTrainingJob(
    display_name="gemma-2b-sft-lora",
    container_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-training-cu121.2-3.transformers.4-42.ubuntu2204.py310",
    command=["trl", "sft"],
)

Runs the Hugging Face PyTorch training DLC with the TRL SFT CLI as the container command.

Configure SFT LoRA arguments

args = [
    "--model_name_or_path=google/gemma-2b",
    "--dataset_name=timdettmers/openassistant-guanaco",
    "--use_peft",
    "--lora_r=16",
    "--lora_alpha=32",
    "--optim=adamw_bnb_8bit",
]

Defines the base model, dataset, LoRA adapter settings, and 8-bit optimizer for memory-efficient fine-tuning.

Submit GPU training job

job.submit(
    args=args,
    replica_count=1,
    machine_type="g2-standard-12",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
    environment_variables={"HF_TOKEN": get_token()},
)

Schedules asynchronous Vertex AI training on one NVIDIA L4 GPU with Hugging Face Hub credentials.

Models & APIs used

  • Models: google/gemma-2b
  • APIs / services: Vertex AI, Cloud Storage, Artifact Registry
  • SDKs / libraries: google-cloud-aiplatform, huggingface_hub, trl, transformers, peft, bitsandbytes

When to use this

Use this pattern to fine-tune a gated Hugging Face Gemma model with TRL SFT and LoRA on managed Vertex AI GPU infrastructure.

Gotchas & caveats

  • google/gemma-2b is gated and requires accepting Google’s usage license on Hugging Face Hub.
  • A Hugging Face user access token is required and passed as HF_TOKEN to the training container.
  • BUCKET_URI must be a valid gs:// bucket because fine-tuned artifacts are written to Cloud Storage.
  • Required IAM roles include Artifact Registry Reader, Vertex AI User, and Storage Object Creator.
  • Required APIs include Vertex AI API, Cloud Storage API, and Artifact Registry API.
  • GCS FUSE uses /gcs/<BUCKET_NAME> inside the container, not gs:// paths.
  • The notebook targets us-central1 by default unless GOOGLE_CLOUD_REGION is set.
  • The example assumes one g2-standard-12 machine with one NVIDIA_L4 GPU and a 3 hour training timeout.

Best practices

  • Use a Hugging Face PyTorch training DLC instead of manually building the training environment.
  • Use LoRA and an 8-bit optimizer to reduce VRAM requirements for fine-tuning.
  • Set output_dir under the /gcs mounted bucket path so trainer outputs are uploaded to Cloud Storage.
  • Pass logging environment variables so container logs stream cleanly to Google Cloud Logs Explorer.
  • Use a read-only or fine-grained Hugging Face token for gated model access.
  • Use a fixed seed value for reproducible training configuration.

Graph View

Backlinks