How to Fine-Tune Pre-Trained Models for Agent Tasks: A Complete TRL Pipeline

Fine-tune pre-trained models for agent tasks using Supervised Fine-Tuning (SFT) with LoRA adapters via the TRL library, leveraging chat-formatted JSONL datasets and the SFTTrainer class to convert generic LLMs into task-oriented agents.

The bojieli/ai-agent-book repository provides a production-ready implementation for adapting large language models into specialized agents. This guide walks through the complete fine-tuning pipeline found in chapter8/prompt-distillation/train_sft_trl.py, covering dataset preparation, LoRA configuration, and training with the Transformer Reinforcement Learning (TRL) library.

Prepare Your Agent Training Dataset

The pipeline expects a JSONL file where each line contains a conversation example. In chapter8/prompt-distillation/train_sft_trl.py, the load_jsonl_dataset() function reads these entries and converts them into a Hugging Face Dataset object.

Each JSON object must contain a messages list following the chat format:

{
  "messages": [
    {"role": "system", "content": "You are a helpful AI assistant."},
    {"role": "user", "content": "Search the web for the latest AI conference."},
    {"role": "assistant", "content": "Here are the top three conferences..."}
  ]
}

This format mirrors the interaction patterns used by agent frameworks like those in chapter1/web-search-agent/agent.py, teaching the model to handle multi-turn conversations and tool-oriented instructions.

Initialize Model Components with Optimized Precision

Before training, the script initializes the tokenizer and base model with specific configurations for memory efficiency. The AutoTokenizer.from_pretrained() call loads the base tokenizer and ensures a pad_token exists, while AutoModelForCausalLM.from_pretrained() loads the model in torch.bfloat16 precision with disabled KV cache for training.

Load Tokenizer and Model

In lines 82-101 of train_sft_trl.py, the initialization follows this pattern:

tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
model.config.use_cache = False

Configure LoRA for Parameter-Efficient Fine-Tuning

When use_lora=True, the script applies Low-Rank Adaptation (LoRA) to reduce memory footprint. The LoraConfig sets rank=32, alpha=16, and targets all linear layers, wrapping the base model with get_peft_model().

This approach preserves the pre-trained weights while adding only 2-4% trainable parameters, making it feasible to fine-tune 30B+ parameter models on consumer hardware.

Set Up Training Parameters and SFTTrainer

The SFTConfig object defines hyperparameters mirroring OpenAI Cookbook defaults: 1 epoch, batch size 4, gradient accumulation steps 4, and learning rate 2e-4 with cosine scheduling.

Chat Template Formatting

The SFTTrainer receives a formatting_func that applies tokenizer.apply_chat_template() to each example's messages. This ensures the model learns the specific chat format required for agent interactions:

def formatting_func(example):
    return tokenizer.apply_chat_template(
        example["messages"], 
        tokenize=False, 
        add_generation_prompt=False
    )

Execute Training and Persist Checkpoints

The trainer.train() call handles distributed training setups (DDP/FSDP) automatically. After training completes, the script saves both the LoRA adapters and tokenizer to the specified output_dir.

Run the complete pipeline via command line:

python chapter8/prompt-distillation/train_sft_trl.py \
  --train_file data/prompt_distillation_lang.jsonl \
  --model_name Qwen/Qwen3-30B-A3B-Instruct-2507 \
  --output_dir models/prompt_distillation_trl \
  --use_lora \
  --lora_rank 32 \
  --lora_alpha 16 \
  --num_train_epochs 1 \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --learning_rate 2e-4

Deploy the Fine-Tuned Agent for Inference

Load the saved adapters using PeftModel.from_pretrained() to merge the LoRA weights with the base model:

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

tokenizer = AutoTokenizer.from_pretrained("models/prompt_distillation_trl")
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-30B-A3B-Instruct-2507",
    device_map="auto",
    torch_dtype="auto"
)
model = PeftModel.from_pretrained(base_model, "models/prompt_distillation_trl")

# Inference

chat = [
    {"role": "system", "content": "You are a tool-calling AI agent."},
    {"role": "user", "content": "What is the weather in Paris today?"},
]
input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt")
output = model.generate(input_ids, max_new_tokens=256)

Summary

• Chat-formatted datasets: Use JSONL files with messages arrays containing system, user, and assistant turns to train agent conversation patterns.
• LoRA configuration: Apply LoraConfig with rank 32 and alpha 16 for efficient fine-tuning of large models without full weight updates.
• TRL SFTTrainer: Leverage this built-in trainer to handle tokenization, batching, and mixed-precision training automatically.
• Inference integration: Load fine-tuned weights via PeftModel.from_pretrained() to deploy task-specific agents capable of tool calling and planning.

Frequently Asked Questions

What dataset format is required for fine-tuning agent tasks?

The bojieli/ai-agent-book implementation requires a JSONL file where each line contains a JSON object with a messages key. This list must include dictionaries with role and content fields representing the conversation history between system, user, and assistant.

Why use LoRA instead of full fine-tuning for agent models?

LoRA (Low-Rank Adaptation) reduces memory requirements by training only low-rank decomposition matrices while freezing base model weights. According to the source code, this allows fine-tuning 30B+ parameter models on limited hardware with minimal performance degradation compared to full fine-tuning.

How does the TRL SFTTrainer handle chat templates?

The SFTTrainer uses a formatting_func parameter that calls tokenizer.apply_chat_template() on each example's messages. This ensures the training data matches the chat format expected by the model, which is critical for agent tasks requiring structured conversation turns.

Can I adjust the learning rate and batch size for different model sizes?

Yes. The SFTConfig in chapter8/prompt-distillation/train_sft_trl.py exposes standard hyperparameters including --learning_rate, --per_device_train_batch_size, and --gradient_accumulation_steps. For larger models, increase gradient accumulation to maintain effective batch sizes while reducing memory usage.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}