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

> Fine-tune pre-trained models for agent tasks using TRL and LoRA adapters. Convert generic LLMs into task-oriented agents with our complete pipeline and chat-formatted datasets. Easy and effective.

- Repository: [Bojie Li/ai-agent-book](https://github.com/bojieli/ai-agent-book)
- Tags: how-to-guide
- Published: 2026-07-30

---

**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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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:

```json
{
  "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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/train_sft_trl.py), the initialization follows this pattern:

```python
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:

```python
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:

```bash
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:

```python
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`](https://github.com/bojieli/ai-agent-book/blob/main/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.