Implementation:Unslothai Unsloth Train On Responses Only

Knowledge Sources	Unsloth TRL Documentation
Domains	NLP, Training
Last Updated	2026-02-07 00:00 GMT

Overview

Wrapper around unsloth_zoo's response masking utility that modifies the trainer's data collator to mask instruction tokens from the loss computation.

Description

train_on_responses_only is a function from unsloth_zoo (re-exported via unsloth/chat_templates.py) that patches an SFTTrainer's data collator to identify instruction and response delimiters in each training sample and set instruction token labels to -100. This ensures only response tokens contribute to the cross-entropy loss.

Usage

Call after creating the SFTTrainer but before calling .train(). The instruction_part and response_part strings must match the exact token sequences used by your chat template to delimit turns.

Code Reference

Source Location

• Repository: unsloth (re-exported from unsloth_zoo)
• File: unsloth/chat_templates.py (L25, L40 re-export)
• Original: unsloth_zoo package

Signature

def train_on_responses_only(
    trainer,
    instruction_part: str,
    response_part: str,
) -> Trainer:
    """
    Modifies trainer to only compute loss on response tokens.

    Args:
        trainer: SFTTrainer instance to patch.
        instruction_part (str): Token delimiter marking start of instruction
            (e.g., "<|start_header_id|>user" for Llama 3).
        response_part (str): Token delimiter marking start of response
            (e.g., "<|start_header_id|>assistant" for Llama 3).

    Returns:
        Modified trainer with patched data collator.
    """

Import

from unsloth.chat_templates import train_on_responses_only

I/O Contract

Inputs

Name	Type	Required	Description
trainer	SFTTrainer	Yes	Configured SFTTrainer instance
instruction_part	str	Yes	Token delimiter for instruction start
response_part	str	Yes	Token delimiter for response start

Outputs

Name	Type	Description
trainer	Trainer	Same trainer with patched data collator that masks instruction tokens

Usage Examples

Llama 3 Response Masking

from unsloth.chat_templates import train_on_responses_only

trainer = UnslothTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=args,
)

# Apply response masking for Llama 3 template
trainer = train_on_responses_only(
    trainer,
    instruction_part="<|start_header_id|>user<|end_header_id|>",
    response_part="<|start_header_id|>assistant<|end_header_id|>",
)

trainer.train()

ChatML Response Masking

trainer = train_on_responses_only(
    trainer,
    instruction_part="<|im_start|>user\n",
    response_part="<|im_start|>assistant\n",
)

Related Pages

Implements Principle

• Principle:Unslothai_Unsloth_Response_Masking

Requires Environment

• Environment:Unslothai_Unsloth_Python_Transformers