Re0-02: HuggingFace TRL SFTTrainer

Hyacehila

The questions in this article can also be addressedRe0-01 : HuggingFace Transformers TrainerRe0-03 : HuggingFace TRL DPOTrainerHow the concept of a relatively close read together is developed in different contexts.

Auxiliary code:sfttrainer.py

I'm going to take it from here.

In the last chapter, use the original. Trainer Run through SFT training. The programme is fully functional, but needs to:

  • Manually calculate the length of the prompt
  • Manual settings for labels (realization of Los Masking)
  • Manual Selection and Configuration DataColator
  • It's a lot of code, it's a lot of error.

This chapter describes the SFTTrainer of the TRL library- A training device for SFT and RL, with a cover reduced handwritten training details.

Learning objectives of this chapter

This chapter covers:

  1. TRL Library What is it and its advantages?
  2. SFTTrainer How to automate Los Masking
  3. Four data formats for SFT and application of the Convention
  4. Flash Attention How to speed up training
  5. chat_template Role and self-defined approach

1. Overview of TRL Library

1.1 What's TRL?

TRL(Transformer Reinforcement Learning) This is the advanced training library provided by Hugging Face:

graph TB
    subgraph Stack["Hugging Face 训练栈"]
        direction TB
        subgraph TRL["TRL (高级封装)"]
            SFTTrainer["SFTTrainer"]
            DPOTrainer["DPOTrainer"]
            PPOTrainer["PPOTrainer 等"]
        end

    subgraph Transformers["Transformers (基础框架)"]
        Trainer["Trainer"]
    end

    TRL -->|基于| Transformers
end

1.2 SFT Trainer vs Trainer Contrast

Functions Native SFTTrainer
Loss Masking Manually Automatic
Chat template application Manual Call Automatic
DataCollor Selection Manual Configuration Automatic
PEFT Integration Manual Call Get peft model Automatic
Data Format Support Pre-treatment required Multiple formats directly support
Volume of code 100+Pre-treatment 10-20 Lines

2. Four data formats supported by SFTT Trainer

This is the focus of understanding SFTT Trainer.

2.1 Format Comparison

Four data formats supported by SFTT Trainer

Format 1: Standard Language Modeling

{"text": "The sky is blue."}
  • Apply scene: Pre-training, renewal of tasks
  • Characteristics: token all calculations of Loss (no Los Masking)

Format 2: Periodic Landscape Modeling Recommendations

{"messages": [
    {"role": "user", "content": "问题"},
    {"role": "assistant", "content": "回答"}
]}
  • Apply scene: Dialogue missions, chat model training
  • Characteristics: Auto-applied bit template, auto-loss Masking

Format 3: Standard Prompt-Complement

{"prompt": "问题", "completion": "回答"}
  • Apply scene: A simple question and answer mission
  • Characteristics: calculation of only part of the command

Format 4: Periodical Prompt-Complement

{"prompt": [{"role": "user", ...}],
 "completion": [{"role": "assistant", ...}]}
  • Characteristics:Prompt-Complement completed with dialogue

2.2 Recommended format: messages (format 2)

For most of the dialogue tasks,Format 2 (messages) The most recommended:

def process_func_simple(example, tokenizer, max_length):
    """简化的预处理函数"""
    messages = [
        {"role": "user", "content": f"请总结: {example['dialogue']}"},
        {"role": "assistant", "content": example['summary']}
    ]
    return {"messages": messages}  # 就这么简单!

Format 3 is acceptable if a simple Prompt-Complement task is required, unless you clearly know your needs, format 1 and 4 are used relatively little, and we will introduce in code the main message of the Conversation Range Modeling, or messageages, while adding a part of the Prompt-Complement.

Compare the previous chapter's procs func

  • Previous chapter: 30+ line code, manual calculation program len
  • This chapter: 5 Line Code, SFTT Trainer Automatically handles everything

3. Internal working principles of SFTT Trainer

3.1 Automated processing processes

When you use the messages format, the SFTTrainer internal will:

Processes

graph TD
    A["messages格式数据<br/>role: user/assistant"] --> B[检测数据格式]
    B --> C[应用 chat_template]
    C --> D[识别 assistant 回复边界]
    D --> E[自动设置 labels]
    E --> F[分词和 padding]
    F --> G[训练就绪的 batch]

Detailed steps

  1. Test Data Format ~ Found in messages format
  2. Apply chat template • Generate standardized dialogue formats
  3. Identify ansistant to border → Automatically recognize by special tag
  4. Autoset labels →prompt part set to -100, ansistant part set token ID
  5. Phrasing and peding • Format available for processing into models

3.2 Key configuration parameters

from trl import SFTConfig

training_args = SFTConfig( output_dir="./output",

# 核心参数:自动化 Loss Masking
assistant_only_loss=True,  # [重点] 只对 assistant 回复计算 Loss

# 可选:自定义 chat_template
chat_template_path=&quot;path/to/template.jinja&quot;,
# 并不是所有模型的chat_template 都原生支持assistant_only_loss=True,
# 这就涉及了自定义chat_template,关于具体的区别,我们会在后面再聊到

# 其他常规参数
max_length=2048,
per_device_train_batch_size=4,
learning_rate=2e-4,
...

)

4. assistant_only_loss vs completion_only_loss

These are two easily confused parameters:

Parameters Applicable Format Role
assistant_only_loss=True Messages (format 2/4) Calculate content for ansistant role only
completion_only_loss=True Prompt-complement (Format 3/4) Calculates only part of the command

Attention.completion_only_loss The default is... True, usually without a visible setting and without an additional setup chat template, the dominant language model on the market is directly supported. On the contrary,assistant_only_loss=TrueDefault asFalse And it needs special chat template support.

4.1 Use messageages format (recommended)

# 需要设置 assistant_only_loss=True
# 可能需要设置 chat_template_path(如果模型默认不支持)
training_args = SFTConfig(
    assistant_only_loss=True,
    chat_template_path="path/to/template.jinja",  # Qwen3 需要
    ...
)

4.2 Use conversional prompt complement

# 不需要设置 assistant_only_loss
# 不需要设置 chat_template_path
# 使用默认配置即可(completion_only_loss=True 是默认值)
training_args = SFTConfig(
    # 默认配置就能工作
    ...
)

Chat template

5.1 What's a chat template?

Chat template is usually a Jinja2 template that defines how to format the dialogue into a text that the model understands:

输入 messages:
[
    {"role": "user", "content": "你好"},
    {"role": "assistant", "content": "你好!有什么可以帮你的吗?"}
]

    ↓ 应用 chat_template ↓

输出文本: <|im_start|>user 你好<|im_end|> <|im_start|>assistant 你好!有什么可以帮你的吗?<|im_end|>

5.2 Why do you need to define a carat template?

Problem: Not all models support return_assistant_tokens_mask Functions

The blogger says that the government is not a party to the law.return_assistant_tokens_mask The new Jinja2 template used for functionality may lead to incompatibility with other frameworks, and therefore, from a future perspective, this fit may well not continue.

This feature enables SFTTrainer to automatically identify the boundaries that are returned by antlert. Customizes are required if the model ' s default template is not supported.

Qwen3

  • Default template not supported generation Keywords
  • Need to use modified template to enable assistant_only_loss

5.3 Custom template example

Customize the template approach to using the components generated by the keywords assistant on the basis of the original template, so that the SFTTrainer can mask them by using the keyword.


{#- 自定义 template 的关键部分 -#}
{%- for message in messages %}
    {%- if message['role'] == 'assistant' %}
        {{- '<|im_start|>assistant\n' }}
        {%- generation %}  {# ← 这个标签告诉 SFTTrainer 这是要训练的部分 #}
            {{- message['content'] + '<|im_end|>' }}
        {%- endgeneration %}
    {%- endif %}
{%- endfor %}

6. Flash Attention: Accelerating training

6.1 What is Flash Act?

Flash Attention is an optimal focus calculation to achieve:

Indicators Standard attention. Flash Attention
Speed Benchmark Come on, 2-3 times.
Organisation Benchmark Savings 50-80 per cent
Long sequence support Limited Longer Sequence

6.2 Enable Flash Attention

# 检查是否可用
try:
    from flash_attn import __version__
    FLASH_ATTENTION_AVAILABLE = True
except ImportError:
    FLASH_ATTENTION_AVAILABLE = False

加载模型时启用

model = AutoModelForCausalLM.from_pretrained( "Qwen/Qwen3-8B", attn_implementation="flash_attention_2", # ← 启用 Flash Attention dtype=torch.bfloat16, … )

6.3 Installation Flash Attention

pip install flash-attn --no-build-isolation

Request

  • GUI supported by CUDA
  • Ampere architecture or updates (A100, H100, RTX 30xx/40xx, etc.)
  • Flash-attn2 supports the Ampere architecture, using flash-attn depending on the support if it has a graphic card to update the architecture

7. Complete code comparison

7.1 Previous Chapter (Previous Trainer)

# 需要手动处理很多事情
def process_func(example, tokenizer, max_length):
    messages = [...]
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    tokenized = tokenizer(text, ...)

# 手动计算 prompt 长度
prompt_text = tokenizer.apply_chat_template(user_only, ...)
prompt_len = len(tokenizer(prompt_text)[&quot;input_ids&quot;])

# 手动设置 labels
labels = [-100] * len(input_ids)
labels[prompt_len:] = input_ids[prompt_len:]

return {&quot;input_ids&quot;: ..., &quot;labels&quot;: labels, ...}

手动选择 DataCollator

data_collator = DataCollatorForSeq2Seq(tokenizer)

手动应用 LoRA

model = get_peft_model(model, lora_config)

trainer = Trainer(model, data_collator=data_collator, …)

7.2 This Chapter (SFTTTrainer)

# 简化的预处理
def process_func_simple(example, tokenizer, max_length):
    messages = [
        {"role": "user", "content": ...},
        {"role": "assistant", "content": ...}
    ]
    return {"messages": messages}  # 就这么简单!

配置 SFTConfig

training_args = SFTConfig( assistant_only_loss=True, # 自动 Loss Masking … )

SFTTrainer 自动处理一切

trainer = SFTTrainer( model=model, args=training_args, peft_config=lora_config, # 自动应用 LoRA train_dataset=train_dataset, # 不需要设置 data_collator! )

8. Operational codes

8.1 Selecting data formats

Yes. sfttrainer.py Bottom changes:

# 选项1: messages 格式(推荐,需要自定义 chat_template)
DATA_FORMAT = "messages"

选项2: conversational_prompt_completion 格式 (因为已经配置好 chat_template)

DATA_FORMAT = "conversational_prompt_completion"

8.2 Operational training

python sfttrainer.py

9. Summary of this chapter

Concept Annotations
SFTTrainer Specialized SFT traininger for TRL library, highly automated
Messages format Recommended dialogue data format, with associated only loss
assistant_only_loss Key parameters for automating Los Masking
chat_template Define formatted dialogue, which may require a self-defined approach in this chapter
Flash Attention Speed up training and save the world.

10. Comparative summary

From Trainer to SFTT Trainer: Comparative summary

Feature Native SFTTrainer
Flexibility High, full control. SFT Optimisation
Loss Masking It needs to be done manually. [Yes] Automatically
DataCollator Manual Configuration Required [Yes] Autoconfiguration
Code Complexity It's a big code. It's a lot of error. Simplicity, best practices

Conclusions Recommended use of SFTTrainer in the production environment of SFT, standard Trainer designed for continuous pre-training

Next chapter.

SFT has taught models to answer questions according to instructions. However:

  • Model answers may not be enough."Okay."
  • Users may have different preferences
  • We want the model's answer to be more in line with human expectations.

Problem: How can models learn about human preferences?

Answer: Use preferred alignment technology!

In the next chapter, we'll learn. DPO(Direct Preference Optimization)

  • Use preference data (chosen vs reprojected)
  • No training incentive model.
  • It's easier to stabilize than the traditional RLHF (PPO)
  • As a bridge to enhanced learning at the senior level

Appendix: Common problems

What if the law doesn't work?
A: Check if the model's chat template supports the model label. Qwen3 needs to use custom template.

Q: Flash Attention installation failed?
A: Ensure that CUDA and PyTorch versions are compatible. Try pip install flash-attn --no-build-isolation

Q: Which of the two data formats should we choose?
A: multi-cycle dialogue options: messages + options only loss; single-cycle questions and answers can be selected for prompt complementation, thus avoiding the question of dialogue templates.

References

  • Title: Re0-02: HuggingFace TRL SFTTrainer
  • Author: Hyacehila
  • Created at : 2025-12-28 14:00:00
  • Link: https://hyacehila.github.io//blog/2025/12/28/Re0HF-02/
  • License: This work is licensed under CC BY-NC-SA 4.0.
Comments