Skip to main content

Fine-tuning

This guide is for software developers interested in fine-tuning Ai2 models (such as Olmo, Tülu, and OlmoE) to better suit their application's specific needs. Fine-tuning is the process of taking a pre-trained language model and further training it on a smaller, task-specific dataset, allowing the model to adapt its general knowledge to your specific use case while maintaining its core capabilities. Think of it as teaching a broadly educated person to specialize in a particular field. This guide provides practical guidance on when fine-tuning is necessary, what resources are required, what fine-tuning data looks like, and how to implement the fine-tuning process.

Try These First​

Before committing to fine-tuning, consider these prompt-engineering alternatives:

  • Basic Prompt Engineering: Best when minor errors are acceptable and rapid iteration is preferred (more info)
  • Few-shot Learning: Few-shot learning involves adding a few examples into your prompt and is effective when you have a small set of high-quality examples to demonstrate the desired behavior (more info)
  • Chain-of-Thought Prompting: Useful for complex reasoning tasks that benefit from step-by-step thinking (more info)

Resource Requirements​

Fine-tuning requires significant resources:

  • Compute: Even fine-tuning a 7B parameter model requires multiple GPUs
  • Data: As a general guideline, a minimum of 1,000 high-quality examples for effective fine-tuning
  • Time: Days to weeks of tuning time, depending on model size, dataset, and number/type of GPUs

Example fine-tuning data formats​

Most Ai2 models expect fine-tuning data in JSONL format, with each line representing a single tuning example.

Instruction-tuning format​

Instruction tuning is a form of supervised fine-tuning (SFT) that teaches models to follow specific instructions or prompts. While SFT broadly refers to any supervised training on labeled examples, instruction tuning specifically focuses on tuning models to understand and execute various types of instructions, from summarization to code generation. This approach helps models become more helpful and aligned with user intentions, making them more effective as AI assistants.

Below is an example of the JSONL format used for instruction tuning and fine-tuning, where each line contains a single JSON object:

{"input": "Summarize the following article:", "output": "This article explains the benefits of solar energy..."}
{"input": "Write a poem about nature:", "output": "The trees sway gently in the breeze, while birds sing melodies with ease..."}

Chat format (for assistant-style models like Tülu)​

Chat format is used for training conversational AI models, where each example contains a sequence of messages between a user and an assistant. This format is particularly important for DPO (Direct Preference Optimization) training, where the model learns from human preferences by comparing different responses to the same input. In DPO, the chat format helps create pairs of preferred and dispreferred responses that teach the model to generate more helpful and aligned responses.

Below is an example of the JSONL format used for DPO training and fine-tuning, where each line contains a single JSON object:

{"messages": [{"role": "user", "content": "How do I reverse a list in Python?"}, {"role": "assistant", "content": "You can use the reverse() method or slicing."}]}
{"messages": [{"role": "user", "content": "What's the difference between a tuple and a list?"}, {"role": "assistant", "content": "Tuples are immutable while lists are mutable. Tuples use parentheses () and lists use square brackets []."}]}

Parameter-efficient fine-tuning (PEFT)​

PEFT is a collection of techniques that enable efficient fine-tuning of large language models by updating only a small subset of parameters. This approach significantly reduces computational requirements and memory usage while maintaining model performance. Common PEFT methods include:

  • LoRA (Low-Rank Adaptation): Adds small, trainable rank decomposition matrices to existing weights
  • Prefix/Prompt Tuning: Prepends trainable vectors to the input or hidden states
  • Adapter Layers: Inserts small trainable layers between existing model layers
  • QLoRA: Combines quantization with LoRA for even greater efficiency

PEFT is particularly valuable when:

  • You have limited computational resources
  • You need to maintain multiple fine-tuned versions of a model
  • You want to quickly experiment with different fine-tuning approaches

For detailed implementation guidance and best practices, see the Hugging Face PEFT documentation.

Fine-tuning example: Olmo 2 32B model​

  1. Create a new Python environment: This step will create an isolated Python environment using venv, which helps manage dependencies and avoid conflicts between different projects. You can create a new Python environment using venv:

    python -m venv olmo_core
    source olmo_core/bin/activate  # On Unix/macOS
    # or
    .\olmo_core\Scripts\activate  # On Windows

  2. Clone and set up Olmo-core: This step will clone the Olmo-core repository and install its dependencies, which provide the necessary components for training the Olmo 2 32B model:

    # Clone OLMo-core repo
    git clone https://github.com/allenai/OLMo-core.git
    cd OLMo-core
    pip install -e ".[all]"
    pip install datasets wandb

  3. Prepare your fine-tuning data: Create a file named data/my-finetune-data.jsonl with the following contents:

    {"input": "Summarize the following article:", "output": "This article explains the benefits of solar energy..."}
    {"input": "Explain the concept of recursion:", "output": "Recursion is a programming concept where a function calls itself..."}

  4. Fine-tune the model: This step will start the fine-tuning process for the Olmo 2 32B model. The script uses distributed fine tuning across multiple GPUs for efficient processing:

    torchrun --nproc-per-node=8 ./src/scripts/official/OLMo-2-0325-32B-train.py finetune01 \
      --model_checkpoint=allenai/OLMo-2-0325-32B \
      --train_dataset=./data/my-finetune-data.jsonl \
      --train_module.optim.lr=3e-6           
      --train_module.global_batch_size=64 \  # Optional: Adjust batch size
      --trainer.max_steps=1000               # Optional: Set maximum training steps

Further reading: model fine tuning and training documentation