Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:Huggingface Trl SFTTrainer Save Model

From Leeroopedia


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 from transformers.Trainer, this saves the model weights (or PEFT adapter weights), tokenizer/processor, and configuration to the specified output_dir.
  • push_to_hub() -- Inherited from transformers.Trainer, this uploads the saved model, tokenizer, configuration, and model card to a HuggingFace Hub repository. The repository name is derived from output_dir or hub_model_id.
  • _save_checkpoint() -- Overridden by SFTTrainer to generate a model card before each checkpoint save. The model card is created via create_model_card() (inherited from BaseTrainer) and documents the base model, dataset, training method, and library versions. The model name is derived from hub_model_id or the output directory name.

In the standard SFT script (trl/scripts/sft.py), the save-and-push sequence is:

  1. trainer.save_model(training_args.output_dir)
  2. If training_args.push_to_hub is 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_checkpoint override)
  • 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)

Related Pages

Implements Principle

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment