SFT Trainer
SFTTrainer performs supervised fine-tuning: it trains a model to imitate the completions in a labelled dataset. It is the first stage of most post-training pipelines.
Dataset format
SFTTrainer accepts two dataset shapes:
- Conversational — a
messagescolumn of{"role", "content"}dicts. The trainer applies the model's chat template automatically. - Standard — a plain
textcolumn, or aprompt/completionpair.
# conversational
{"messages": [
{"role": "user", "content": "What is the capital of France?"},
{"role": "assistant", "content": "The capital of France is Paris."},
]}
Full example with LoRA
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
dataset = load_dataset("trl-lib/Capybara", split="train")
trainer = SFTTrainer(
model="meta-llama/Llama-3.2-1B",
train_dataset=dataset,
args=SFTConfig(
output_dir="Llama-3.2-1B-SFT",
num_train_epochs=1,
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
learning_rate=2e-4,
bf16=True,
packing=True,
),
peft_config=LoraConfig(r=16, lora_alpha=32, task_type="CAUSAL_LM"),
)
trainer.train()
Key SFTConfig arguments
| Argument | Purpose |
|---|---|
packing | Concatenate short samples into one sequence for throughput. |
max_length | Truncate/pack sequences to this many tokens. |
completion_only_loss | Mask the prompt so loss is computed only over the completion. |
assistant_only_loss | For conversational data, train only on assistant turns. |