Skip to main content
Inferix

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 messages column of {"role", "content"} dicts. The trainer applies the model's chat template automatically.
  • Standard — a plain text column, or a prompt/completion pair.
# 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

ArgumentPurpose
packingConcatenate short samples into one sequence for throughput.
max_lengthTruncate/pack sequences to this many tokens.
completion_only_lossMask the prompt so loss is computed only over the completion.
assistant_only_lossFor conversational data, train only on assistant turns.

    We use cookies for essential functionality and analytics. You can accept or reject analytics cookies.Cookie policy