Supervised Fine-Tuning (SFT)

“SFT is a post-training technique used to align LLM responses to a set of desired responses using a dataset of (prompt, response) pairs. This technique is the simplest way to shift a model’s behaviour toward a target style or task and does not utilise reinforcement learning.”

It’s similar to a continuation of the pre-training stage of an LLM, but using a curated dataset that is specific to the LLM’s application. Cross-entropy loss is computed exclusively on the response tokens, so the model is never penalised for how it encodes the prompt.

SFT is typically the first stage of a two-step alignment pipeline:

  • SFT (this class): warm-up the model to follow instructions by minimising cross-entropy on (prompt, good_response) pairs.

  • DPO: further align the SFT-initialised model using (prompt, chosen_response, rejected_response) triples.

This technique is surprisingly effective, as pre-trained LLMs have been shown to easily adapt to a relatively small amount of new data.

Example

from agilerl.algorithms.sft import SFT
from agilerl.llm_envs import DatasetEnv
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Instantiate the model and the associated tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B",
    torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")

# Load the dataset into an SFT DatasetEnv
raw_dataset = load_dataset("HumanLLMs/Human-Like-DPO-Dataset", split="train").shuffle(seed=42)
train_test_split = raw_dataset.train_test_split(test_size=0.1)
train_dataset = train_test_split["train"]
test_dataset = train_test_split["test"]
env = DatasetEnv(
  train_dataset=train_dataset,
  test_dataset=test_dataset,
  tokenizer=tokenizer,
  objective="sft",
  response_column="chosen",
  data_batch_size_per_gpu=16,
)

# Configure LoRA adapters
lora_config = LoraConfig(
  r=16,
  lora_alpha=64,
  target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
  lora_dropout=0.05,
  bias="none",
)

# Instantiate the agent
agent = SFT(
  actor_network=model,
  pad_token_id=tokenizer.eos_token_id,
  pad_token=tokenizer.eos_token,
  device="cuda" if torch.cuda.is_available() else "cpu",
  batch_size=32,
  lr=5e-5,
  update_epochs=1,
  lora_config=lora_config,
  seed=42,
)

Training an SFT agent

To train an SFT agent on a single dataset environment, use the train_llm_dataset function:

from agilerl.training.llm import train_llm_dataset

train_llm_dataset(
  pop=[agent],
  env=env,
  init_hp={"BATCH_SIZE": 32, "UPDATE_EPOCHS": 1},
  checkpoint_steps=250,
)

Saving and Loading Agents

To save an agent, use the save_llm_checkpoint function:

from agilerl.utils.utils import save_llm_checkpoint

save_llm_checkpoint(agent, "path/to/checkpoint")

To load a trained model, you must use the HuggingFace .from_pretrained method, AgileRL is compatible with HuggingFace and Peft models:

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
model = PeftModel.from_pretrained(base_model, "path/to/model/directory")

Parameters