Generative AI Google Brain

open-models-fine-tuning-get-started-with-oss-tuning-on-vertexai

3 min read

MetaMath with Vertex AI Open Source Model Tuning

Source notebook

Repo path: open-models/fine-tuning/get_started_with_oss_tuning_on_vertexai.ipynb · Open on GitHub · advanced

Fine-tunes a Llama 3.1 8B model on MetaMathQA using Vertex AI managed tuning.

Summary

This notebook shows how to prepare the MetaMathQA dataset for Vertex AI supervised fine-tuning, upload JSONL training and validation files to Cloud Storage, and launch a full tuning job for an open-source Llama model. It then deploys the tuned model with Vertex AI Model Garden, sends a math reasoning prompt to the endpoint, optionally compares against the official Hugging Face MetaMath model, and optionally runs the MetaMath GSM8K evaluation script.

Key code patterns

Initialize Vertex AI

PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
LOCATION = "europe-west4"
BUCKET_URI = f"gs://{BUCKET_NAME}"
! gcloud storage buckets create {BUCKET_URI} --project {PROJECT_ID} --location {LOCATION}
vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)

Sets project, region, and staging bucket before creating tuning and deployment resources.

Format MetaMathQA for tuning

def format_for_tuning(example):
    instruction = METAMATH_TEMPLATE.format(instruction=example["query"])
    return {"messages": [
        {"role": "user", "content": instruction},
        {"role": "assistant", "content": f" {example['response']}"},
    ]}
train_formatted_dataset = train_split.map(format_for_tuning, remove_columns=train_split.column_names)

Converts Hugging Face dataset rows into the chat-style JSONL structure expected by Vertex AI tuning.

Launch managed SFT job

source_model = SourceModel(base_model=config.base_model)
sft_tuning_job = sft.train(
    source_model=source_model,
    tuning_mode=config.tuning_mode,
    epochs=config.epochs,
    learning_rate=config.learning_rate,
    train_dataset=train_file_uri,
    validation_dataset=validation_file_uri,
    output_uri=output_uri,
)

Uses Vertex AI managed supervised fine-tuning with explicit base model, FULL tuning mode, datasets, and output URI.

Deploy tuned artifacts

tuned_model = model_garden.CustomModel(gcs_uri=model_artifacts_gcs_uri)
endpoint = tuned_model.deploy(
    machine_type="g2-standard-12",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
)

Turns tuned GCS model artifacts into a deployed Vertex AI endpoint for inference.

Predict with tuned endpoint

instances = [{
    "prompt": prompt_template.format(instruction=instruction),
    "max_tokens": 250,
    "temperature": 0.2,
    "top_p": 1.0,
    "top_k": 1,
    "raw_response": True,
}]
response = endpoint.predict(instances=instances, use_dedicated_endpoint=True)

Uses the MetaMath prompt template and deterministic decoding settings to test math reasoning.

Models & APIs used

  • Models: meta/llama3_1@llama-3.1-8b, meta-math/MetaMath-7B-V1.0
  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, vertexai, datasets, transformers, torch, pydantic

When to use this

Use this pattern to reproduce or adapt paper-style open-model fine-tuning on Vertex AI with managed training, deployment, and benchmark comparison.

Gotchas & caveats

  • Vertex AI API must be enabled for the Google Cloud project.
  • The notebook creates and uses a Cloud Storage bucket for datasets and model artifacts.
  • Validation data is capped below 5000 rows for a Vertex AI requirement.
  • Managed tuning can take several hours, and endpoint deployment can take 15-30 minutes.
  • Optional local comparison with MetaMath-7B-V1.0 may require significant RAM, GPU, download time, and a larger Workbench runtime.
  • Official evaluation uses tensor_parallel_size 2 and suggests an A100-based Workbench instance for resource-heavy sections.
  • Cleanup is needed to avoid ongoing endpoint and storage charges.

Best practices

  • Use an 80/20 train-validation split with a fixed seed for reproducibility.
  • Limit validation rows to satisfy the Vertex AI validation dataset requirement.
  • Upload JSONL files to Cloud Storage because the Vertex AI tuning service cannot access local notebook files directly.
  • Use the same MetaMath prompt template for tuned-model testing and official-model comparison.
  • Use low temperature, top_p 1.0, and top_k 1 for factual math output.
  • Monitor long-running tuning jobs by refreshing job state periodically.
  • Delete experiments, endpoints, and Cloud Storage artifacts during cleanup.

Graph View

Backlinks