Fine-tuning with FunctionGemma

View on ai.google.dev Run in Google Colab Run in Kaggle Open in Vertex AI View source on GitHub

This guide demonstrates how to fine-tune FunctionGemma for tool calling.

While FunctionGemma is natively capable of calling tools, true capability comes from two distinct skills: the mechanical knowledge of how to use a tool (syntax) and the cognitive ability to interpret why and when to use it (intent).

Models, especially smaller ones, have fewer parameters available to retain complex intent understanding. This is why we need to fine-tune them.

Common use cases for fine-tuning tool calling include:

  • Model Distillation: Generating synthetic training data with a larger model and fine-tuning a smaller model to replicate the specific workflow efficiently.
  • Handling Non-Standard Schemas: Overcoming base model struggles with legacy, highly complex data structures or proprietary format not found in public data, such as handling domain-specific mobile actions.
  • Optimizing Context Usage: "Baking" tool definitions into the model's weights. This allows you to use shorthand descriptions in your prompts, freeing up the context window for the actual conversation.
  • Resolving Selection Ambiguity: Biasing the model toward specific enterprise policies, such as prioritizing an internal knowledge base over an external search engine.

In this example, we will focus specifically on managing tool selection ambiguity.

Setup development environment

The first step is to install Hugging Face Libraries, including TRL, and datasets to fine-tune open model, including different RLHF and alignment techniques.

# Install Pytorch & other libraries
%pip install torch tensorboard

# Install Hugging Face libraries
%pip install transformers datasets accelerate evaluate trl protobuf sentencepiece

# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100
#% pip install flash-attn

Note: If you are using a GPU with Ampere architecture (such as NVIDIA L4) or newer, you can use Flash attention. Flash Attention is a method that significantly speeds computations up and reduces memory usage from quadratic to linear in sequence length, leading to accelerating training up to 3x. Learn more at FlashAttention.

Before you can start training, you have to make sure that you accepted the terms of use for Gemma. You can accept the license on Hugging Face by clicking on the Agree and access repository button on the model page at: http://huggingface.co/google/functiongemma-270m-it

After you have accepted the license, you need a valid Hugging Face Token to access the model. If you are running inside a Google Colab, you can securely use your Hugging Face Token using the Colab secrets otherwise you can set the token as directly in the login method. Make sure your token has write access too, as you push your model to Hugging Face Hub after fine-tuning.

# Login into Hugging Face Hub
from huggingface_hub import login
login()

You can keep the results on Colab's local virtual machine. However, it is highly recommended saving your intermediate results to your Google Drive. This ensures your training results are safe and allows you to easily compare and select the best model.

Also, adjust the checkpoint directory and the learning rate.

from google.colab import drive

mount_google_drive = False
checkpoint_dir = "functiongemma-270m-it-simple-tool-calling"

if mount_google_drive:
    drive.mount('/content/drive')
    checkpoint_dir = f"/content/drive/MyDrive/{checkpoint_dir}"

print(f"Checkpoints will be saved to {checkpoint_dir}")

base_model = "google/functiongemma-270m-it"
learning_rate = 5e-5

Checkpoints will be saved to functiongemma-270m-it-simple-tool-calling

Prepare the fine-tuning dataset

You will use the following example dataset, which contains sample conversations requiring a choice between two tools: search_knowledge_base and search_google.

Simple tool calling dataset