Implementation:Huggingface Trl SFTTrainer Save Model
| Knowledge Sources | |
|---|---|
| Domains | NLP, Training |
| Last Updated | 2026-02-06 17:00 GMT |
Overview
Concrete tools for saving the fine-tuned model to disk and publishing it to the HuggingFace Hub, including model card generation, provided by the TRL library (extending transformers.Trainer).
Description
The SFTTrainer provides model saving through inherited and overridden methods:
save_model()-- Inherited fromtransformers.Trainer, this saves the model weights (or PEFT adapter weights), tokenizer/processor, and configuration to the specifiedoutput_dir.
push_to_hub()-- Inherited fromtransformers.Trainer, this uploads the saved model, tokenizer, configuration, and model card to a HuggingFace Hub repository. The repository name is derived fromoutput_dirorhub_model_id.
_save_checkpoint()-- Overridden bySFTTrainerto generate a model card before each checkpoint save. The model card is created viacreate_model_card()(inherited fromBaseTrainer) and documents the base model, dataset, training method, and library versions. The model name is derived fromhub_model_idor the output directory name.
In the standard SFT script (trl/scripts/sft.py), the save-and-push sequence is:
trainer.save_model(training_args.output_dir)- If
training_args.push_to_hubis True:trainer.push_to_hub(dataset_name=script_args.dataset_name)
Usage
Use save_model() after training completes to persist the model. Use push_to_hub() to distribute the model via the HuggingFace Hub. Checkpoints during training are handled automatically based on the save_strategy configuration.
Code Reference
Source Location
- Repository: TRL
- File:
trl/trainer/sft_trainer.py(lines 1306-1312,_save_checkpointoverride) - File:
trl/scripts/sft.py(lines 149-155, save and push sequence) - File:
trl/trainer/utils.py(lines 434-509,generate_model_card)
Signature
class SFTTrainer(BaseTrainer):
# Inherited from transformers.Trainer
def save_model(self, output_dir: str | None = None) -> None:
"""
Save the model, tokenizer, and configuration to output_dir.
For PEFT models, saves only the adapter weights.
"""
...
# Inherited from transformers.Trainer
def push_to_hub(self, **kwargs) -> str:
"""
Upload model to HuggingFace Hub.
Returns the URL of the Hub repository.
"""
...
# SFT-specific override
def _save_checkpoint(self, model, trial):
"""
Generate model card before saving checkpoint.
Derives model name from hub_model_id or output_dir.
"""
if self.args.hub_model_id is None:
model_name = Path(self.args.output_dir).name
else:
model_name = self.args.hub_model_id.split("/")[-1]
self.create_model_card(model_name=model_name)
super()._save_checkpoint(model, trial)
Import
# Methods are called on an SFTTrainer instance; no separate import needed
from trl import SFTTrainer
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| output_dir | None | No | Directory to save the model to; defaults to args.output_dir
|
| push_to_hub kwargs | dict |
No | Additional arguments for Hub upload, such as dataset_name for the model card
|
Outputs
| Name | Type | Description |
|---|---|---|
| Saved files (disk) | files | Model weights (model.safetensors or adapter files), config.json, tokenizer files, README.md (model card)
|
| Hub URL (push_to_hub) | str |
URL of the uploaded model on the HuggingFace Hub |
Usage Examples
Save Model to Disk
from trl import SFTTrainer, SFTConfig
from datasets import load_dataset
dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]")
trainer = SFTTrainer(
model="Qwen/Qwen2.5-0.5B-Instruct",
args=SFTConfig(output_dir="./my-sft-model"),
train_dataset=dataset,
)
trainer.train()
# Save the final model
trainer.save_model("./my-sft-model")
# Output directory contents:
# ./my-sft-model/
# config.json
# model.safetensors (or adapter_model.safetensors for PEFT)
# tokenizer.json
# tokenizer_config.json
# special_tokens_map.json
# README.md (model card)
Save and Push to Hub
trainer = SFTTrainer(
model="Qwen/Qwen2.5-0.5B-Instruct",
args=SFTConfig(
output_dir="./my-sft-model",
push_to_hub=True,
hub_model_id="my-username/Qwen2.5-0.5B-SFT",
),
train_dataset=dataset,
)
trainer.train()
# Save locally
trainer.save_model("./my-sft-model")
# Push to HuggingFace Hub
trainer.push_to_hub(dataset_name="trl-lib/Capybara")
# Model available at: https://huggingface.co/my-username/Qwen2.5-0.5B-SFT
Complete SFT Script Pattern
# Pattern from trl/scripts/sft.py
trainer.train()
trainer.save_model(training_args.output_dir)
if training_args.push_to_hub:
trainer.push_to_hub(dataset_name=script_args.dataset_name)