How to intermediate · 3 min read

How to use SFTTrainer in Hugging Face TRL

Q: How to use SFTTrainer in Hugging Face TRL

Use SFTTrainer from Hugging Face's trl library to fine-tune language models with supervised data. Initialize SFTTrainer with your model, tokenizer, and dataset, then call train() to start fine-tuning.

Quick answer

Use SFTTrainer from Hugging Face's trl library to fine-tune language models with supervised data. Initialize SFTTrainer with your model, tokenizer, and dataset, then call train() to start fine-tuning.

PREREQUISITES

Python 3.8+
pip install transformers>=4.30.0
pip install trl>=0.4.0
Access to a pretrained Hugging Face model
Basic knowledge of PyTorch or TensorFlow

Setup

Install the required libraries and set up your environment variables if needed. Ensure you have transformers and trl installed.

bash

pip install transformers>=4.30.0 trl>=0.4.0

Step by step

This example demonstrates fine-tuning a GPT-2 model using SFTTrainer with a simple supervised dataset.

python

from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer
from datasets import Dataset

# Load pretrained model and tokenizer
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Prepare a simple dataset
train_data = [{"input_text": "Hello, how are you?", "target_text": "I am fine, thank you!"},
              {"input_text": "What is AI?", "target_text": "AI stands for Artificial Intelligence."}]
dataset = Dataset.from_list(train_data)

# Tokenize function
def tokenize_function(examples):
    inputs = tokenizer(examples["input_text"], truncation=True, padding="max_length", max_length=32)
    targets = tokenizer(examples["target_text"], truncation=True, padding="max_length", max_length=32)
    inputs["labels"] = targets["input_ids"]
    return inputs

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# Initialize SFTTrainer
trainer = SFTTrainer(
    model=model,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    max_length=32,
    batch_size=2,
    max_epochs=3,
    learning_rate=5e-5
)

# Start training
trainer.train()

output

***** Running training *****
  Num examples = 2
  Num Epochs = 3
  Instantaneous batch size per device = 2
...
Training completed.

Common variations

Use different pretrained models by changing model_name (e.g., gpt2-medium, EleutherAI/gpt-j-6B).
Adjust max_length, batch_size, and learning_rate for your dataset size and hardware.
Use SFTTrainer with custom datasets loaded from files or Hugging Face Hub.
For asynchronous or distributed training, integrate with PyTorch Lightning or accelerate.

Troubleshooting

If you encounter tokenization errors, verify your tokenizer and dataset inputs match in length and format.
Out of memory errors? Reduce batch_size or max_length.
Training does not start? Confirm your dataset is properly tokenized and contains labels.
Check that your transformers and trl versions are compatible.

✅

Key Takeaways

Use SFTTrainer to fine-tune Hugging Face causal language models with supervised datasets.
Tokenize inputs and targets properly, setting labels for supervised fine-tuning.
Adjust training parameters like batch size and learning rate based on your hardware and dataset.
Ensure compatible versions of transformers and trl to avoid runtime errors.

Verified 2026-04 · gpt2, gpt2-medium, EleutherAI/gpt-j-6B

Verify ↗