TRL SFTTrainer on a Dedicated GPU

Axolotl, Unsloth, and every other fine-tuning framework ultimately wrap Hugging Face’s TRL SFTTrainer. Using it directly on dedicated GPU hosting gives you full visibility and control. Here is a minimal complete example.

Contents

• Install
• Dataset format
• Training script
• Common tuning

Install

pip install torch transformers trl peft datasets accelerate bitsandbytes

Dataset

SFTTrainer accepts a Hugging Face Dataset with a messages column (chat format) or a plain text column. Chat format is easier for modern instruct tuning:

[
  {"messages": [
    {"role": "user", "content": "..."},
    {"role": "assistant", "content": "..."}
  ]},
  ...
]

Script

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig

ds = load_dataset("json", data_files="train.jsonl", split="train")
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype="bfloat16",
    device_map="cuda",
)
lora = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj","k_proj","v_proj","o_proj"])

trainer = SFTTrainer(
    model=model,
    args=SFTConfig(
        output_dir="./out",
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        learning_rate=1e-4,
        bf16=True,
        num_train_epochs=3,
        max_seq_length=2048,
    ),
    train_dataset=ds,
    peft_config=lora,
    tokenizer=tok,
)
trainer.train()
trainer.save_model("./out")

Tuning

• max_seq_length: match your dataset’s p95 length; longer wastes compute on padding
• packing=True in SFTConfig: batches multiple short samples into one sequence – 20-40% throughput gain
• gradient_checkpointing=True: trade compute for VRAM when batch size is capped

Compare Axolotl (config-driven) and Unsloth (faster kernels).