Implementation:Hiyouga LLaMA Factory SFT Trainer

Knowledge Sources	Hiyouga_LLaMA_Factory
Domains	Supervised Fine-Tuning, Seq2Seq Training, Trainer
Last Updated	2026-02-06 19:00 GMT

Overview

CustomSeq2SeqTrainer is a custom HuggingFace Seq2SeqTrainer subclass for supervised fine-tuning with support for FP8, custom loss functions, and generative prediction.

Description

CustomSeq2SeqTrainer extends Seq2SeqTrainer with several enhancements: FP8 environment configuration and verification at initialization, custom optimizer and scheduler creation (supporting GaLore, BAdam, and other advanced optimizers), optional DFT and EAFT loss functions, and BAdam callback integration. The prediction_step override strips prompt tokens from generated outputs by replacing them with pad tokens. The save_predictions method writes prompt/predict/label triples as JSONL. It also handles correct loss computation under gradient accumulation for multimodal processors.

Usage

Use CustomSeq2SeqTrainer for supervised fine-tuning of language models with both standard cross-entropy training and generative evaluation/prediction. It is instantiated by the run_sft workflow function and supports both text-only and multimodal inputs.

Code Reference

Source Location

• Repository: Hiyouga_LLaMA_Factory
• File: src/llamafactory/train/sft/trainer.py
• Lines: 1-184

Signature

class CustomSeq2SeqTrainer(Seq2SeqTrainer):
    def __init__(
        self,
        finetuning_args: "FinetuningArguments",
        processor: Optional["ProcessorMixin"],
        model_args: Optional["ModelArguments"] = None,
        gen_kwargs: Optional[dict[str, Any]] = None,
        **kwargs,
    ) -> None

    def create_optimizer(self) -> "torch.optim.Optimizer"

    def create_scheduler(
        self,
        num_training_steps: int,
        optimizer: Optional["torch.optim.Optimizer"] = None,
    ) -> "torch.optim.lr_scheduler.LRScheduler"

    def prediction_step(
        self,
        model: "torch.nn.Module",
        inputs: dict[str, Union["torch.Tensor", Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[list[str]] = None,
        **gen_kwargs,
    ) -> tuple[Optional[float], Optional["torch.Tensor"], Optional["torch.Tensor"]]

    def save_predictions(
        self,
        dataset: "Dataset",
        predict_results: "PredictionOutput",
        skip_special_tokens: bool = True,
    ) -> None

Import

from llamafactory.train.sft.trainer import CustomSeq2SeqTrainer

I/O Contract

Inputs

Name	Type	Required	Description
finetuning_args	FinetuningArguments	Yes	Fine-tuning configuration including use_badam, use_dft_loss, use_eaft_loss, disable_shuffling, and compute_accuracy
processor	Optional[ProcessorMixin]	Yes	Multimodal processor; triggers SaveProcessorCallback and disables model_accepts_loss_kwargs
model_args	Optional[ModelArguments]	No	Model arguments (reserved for future use)
gen_kwargs	Optional[dict[str, Any]]	No	Generation keyword arguments for predict_with_generate mode
**kwargs	dict	Yes	Passed to parent Seq2SeqTrainer; must include model, args, tokenizer, data_collator, etc.

Outputs

Name	Type	Description
prediction_step result	tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]	(loss, generated_tokens_with_prompt_stripped, labels)
generated_predictions.jsonl (from save_predictions)	File	JSONL file with {"prompt": ..., "predict": ..., "label": ...} per example

Usage Examples

# Typically instantiated by run_sft, not directly
from llamafactory.train.sft.trainer import CustomSeq2SeqTrainer

trainer = CustomSeq2SeqTrainer(
    model=model,
    args=training_args,
    finetuning_args=finetuning_args,
    data_collator=data_collator,
    callbacks=callbacks,
    gen_kwargs=gen_kwargs,
    processor=processor,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=ComputeSimilarity(tokenizer=tokenizer),
)

# Training
train_result = trainer.train()

# Prediction with output saving
predict_results = trainer.predict(eval_dataset)
trainer.save_predictions(eval_dataset, predict_results)
# Output JSONL: {"prompt": "...", "predict": "...", "label": "..."}

Related Pages

• Hiyouga_LLaMA_Factory_SFT_Workflow - The workflow orchestrator that creates and drives CustomSeq2SeqTrainer
• Hiyouga_LLaMA_Factory_SFT_Metric - Metric classes (ComputeAccuracy, ComputeSimilarity) used with this trainer
• Hiyouga_LLaMA_Factory_FP8_Utils - FP8 environment configuration utilities used at initialization
• Hiyouga_LLaMA_Factory_Trainer_Utils - Custom optimizer and scheduler creation functions