Principle:Huggingface Trl SFT Dataset Preparation
| Knowledge Sources | |
|---|---|
| Domains | NLP, Training |
| Last Updated | 2026-02-06 17:00 GMT |
Overview
Loading, mixing, and formatting datasets into tokenized sequences suitable for supervised fine-tuning, including chat template application and multi-dataset composition.
Description
The dataset preparation stage is the bridge between raw training data and the tokenized tensors that the model consumes. For supervised fine-tuning, this stage must handle several challenges:
- Multiple data formats -- SFT datasets come in three primary formats:
- Language modeling (
{"messages": [...]}or{"text": "..."}): The model is trained on the full sequence. - Prompt-completion (
{"prompt": "...", "completion": "..."}): The model is optionally trained only on the completion portion. - Conversational (messages with
"role"and"content"fields): Multi-turn conversations that require chat template rendering.
- Language modeling (
- Chat template application -- Conversational data must be rendered through the tokenizer's chat template (a Jinja2 template) to produce the special tokens, role markers, and formatting that the model expects. The
apply_chat_template()function handles this conversion, supporting both language modeling and prompt-completion variants.
- Dataset mixing -- Production SFT often trains on a mixture of datasets (e.g., instruction-following + coding + math). The
DatasetMixtureConfigandget_dataset()system allows combining multiple HuggingFace datasets by concatenation, with optional train/test splitting.
- Format normalization -- Legacy conversational datasets use
"from"/"value"keys instead of"role"/"content". The pipeline auto-detects and converts these viamaybe_convert_to_chatml().
- EOS token handling -- For non-conversational text, the pipeline appends the EOS token to each example to ensure the model learns to terminate generation.
Usage
Use this pattern when:
- Loading datasets from the HuggingFace Hub for SFT training.
- Mixing multiple datasets with different schemas into a unified training set.
- Converting conversational data into the model's expected chat format.
- Preparing prompt-completion data with proper loss masking boundaries.
Theoretical Basis
Instruction Fine-Tuning: The standard SFT objective trains the model to maximize the log-likelihood of the target text given context:
L = -sum_{t} log P(y_t | y_{<t}, x)
where x is the prompt/instruction and y is the target completion.
Completion-Only Loss: When completion_only_loss=True, the loss is masked to exclude prompt tokens:
L = -sum_{t in completion} log P(y_t | y_{<t}, x)
This prevents the model from wasting capacity on reproducing the prompt, focusing gradient signal entirely on generating correct completions.
Assistant-Only Loss: For multi-turn conversations, assistant_only_loss=True further restricts the loss to only assistant turns, masking system prompts, user messages, and formatting tokens.
Chat Templates: A chat template T is a Jinja2 function that maps a list of messages to a formatted string:
T(messages) -> "<|system|>\n{system}\n<|user|>\n{user}\n<|assistant|>\n{assistant}<|end|>"
The specific tokens and format are model-dependent, making the template a critical part of data preparation.
Dataset Mixing: Given datasets D_1, D_2, ..., D_n, the mixture is:
D_mix = concatenate(D_1, D_2, ..., D_n)
Optionally split into train/test: D_mix -> (D_train, D_test) based on test_split_size.