Teaching a model function calling with LoRA fine-tuning

This guide is for readers who already know the basic agent loop and want to understand how a model learns to call tools. By the end you will know how function calls appear in a conversation, which training steps a base model goes through before it is useful, and why LoRA can fine-tune a large model with modest resources.

What function calling is

Function calling is a way for an LLM to take actions on its environment. It was first introduced in GPT-4 and was later reproduced in other models^[1].

It is similar to the tools an agent uses, because both give the model the ability to act on its environment. The difference is that function calling is a capability the model learns, and it relies less on prompting than other agent techniques^[1].

In Unit 1, the agent did not learn to use tools. The tools were listed for it, and the model generalised to plan how to use them. With function calling, the agent is fine-tuned, meaning trained, to use tools^[1].

How the model acts: think, act, observe

A model that uses tools goes through a cycle of three steps^[1]:

1. Think: decide which actions are needed to fulfil the objective.
2. Act: format the action with the correct parameters, then stop generating.
3. Observe: get back the result of the execution.

Through an API, a conversation alternates between user and assistant messages, like this^[1]:

conversation = [
    {"role": "user", "content": "I need help with my order"},
    {"role": "assistant", "content": "I'd be happy to help. Could you provide your order number?"},
    {"role": "user", "content": "It's ORDER-123"},
]

Function calling adds two new roles to the conversation: one for an action and one for an observation^[1]. The Mistral API example below shows the assistant making a function call, followed by a message with the tool role that carries the result^[1]:

conversation = [
    {
        "role": "user",
        "content": "What's the status of my transaction T1001?"
    },
    {
        "role": "assistant",
        "content": "",
        "function_call": {
            "name": "retrieve_payment_status",
            "arguments": "{\"transaction_id\": \"T1001\"}"
        }
    },
    {
        "role": "tool",
        "name": "retrieve_payment_status",
        "content": "{\"status\": \"Paid\"}"
    },
    {
        "role": "assistant",
        "content": "Your transaction T1001 has been successfully paid."
    }
]

Many APIs return the action as an assistant message. The chat template then represents it with special tokens for function calling^[1]. They are [AVAILABLE_TOOLS] to start the list of available tools and [/AVAILABLE_TOOLS] to end it, [TOOL_CALLS] to call a tool, which is the action, [TOOL_RESULTS] to observe the result, and [/TOOL_RESULTS] to end the observation, after which the model can decode again^[1]. To go deeper, the Mistral documentation on function calling is a good reference^[1].

Three training steps

A model usually goes through three training steps^[2]:

1. Pre-training on a large quantity of data produces a base model. It only knows how to predict the next token and does not follow instructions well.
2. Fine-tuning teaches it to follow instructions, which makes it useful in a chat context.
3. Alignment adjusts it to the creator's preferences. A customer service chat model, for example, must never be impolite to customers.

A complete product such as Gemini or Mistral usually goes through all three steps. Models on Hugging Face have completed one or more of them^[2].

Why start from an instruction-tuned model

This guide builds on google/gemma-2-2b-it, not on the base model google/gemma-2-2b. The fine-tuned model has been improved for this kind of use case^[2].

Starting from the pre-trained model would require more training to learn instruction following, chat and function calling all at once. Starting from the instruction-tuned model means the model has much less to learn^[2].

LoRA: train only a small part of the model

LoRA (Low-Rank Adaptation of Large Language Models) is a popular, lightweight training technique. It significantly reduces the number of trainable parameters^[2].

It inserts a smaller number of new weights into the model as an adapter, and only the adapter is trained. This makes training faster and more memory-efficient, and it produces small weights, a few hundred MBs, that are easier to store and share^[2]. Concretely, LoRA adds pairs of rank decomposition matrices to Transformer layers, typically focusing on linear layers. The rest of the model is frozen, and only the adapter weights are updated^[2].

During inference, the input passes through both the adapter and the base model. The adapter weights can also be merged into the base model, which adds no extra latency^[2]. LoRA is especially useful for adapting large language models to specific tasks or domains while keeping resource requirements manageable^[2].

Setup: open the notebook

A notebook that runs in a Colab Notebook is available, and you follow its steps from there^[2].

Next steps

The Unit 1 material says the next step is to add function-calling capabilities to google/gemma-2-2b-it, a model that does not have them yet, by appending new special tokens to the model^[1].

Adapted from bonus unit 1 of Hugging Face's open-source course *Agents Course*^[1][2], released under the Apache-2.0 License.