Data Collators: Padding and Beyond

~22 min · training, collator

Level 0Scout

0 XP0/50 lessons0/10 achievements

0/120 XP to next level120 XP to go0% complete

The collator's job

A data collator takes a list of tokenized examples (ragged lengths) and produces a single batch tensor (uniform shape). The standard ones:

DataCollatorWithPadding — pads input_ids to the longest in the batch.
DataCollatorForLanguageModeling — for MLM (masked) or CLM (causal) pre-training. mlm=True randomly masks tokens.
DataCollatorForSeq2Seq — pads encoder + decoder side independently for encoder-decoder models.
DataCollatorForCompletionOnlyLM (from trl) — masks the prompt portion so loss only flows over the assistant response.

Why collators matter for SFT

If you train on chat data without a completion-only collator, the model is also learning to predict the user prompt. That's wasted gradient and often degrades chat behavior. The completion-only collator is the right default for instruction tuning.

Code

Padding collator (classification)·python

from transformers import DataCollatorWithPadding
collator = DataCollatorWithPadding(tokenizer=tok, padding="longest")  # or "max_length"

Completion-only collator (instruction tuning)·python

from trl import DataCollatorForCompletionOnlyLM

# Mark the assistant header so loss only flows after it
response_template = "<|start_header_id|>assistant<|end_header_id|>\n\n"
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template,
    tokenizer=tok,
)

MLM collator (BERT-style pre-training)·python

from transformers import DataCollatorForLanguageModeling
collator = DataCollatorForLanguageModeling(
    tokenizer=tok,
    mlm=True,
    mlm_probability=0.15,  # 15% masking, the BERT default
)

External links

Exercise

Take a small instruction-tuning dataset (e.g., HuggingFaceH4/no_robots). Tokenize it with the chat template applied. Set up two trainers: one with DataCollatorWithPadding, one with DataCollatorForCompletionOnlyLM. Train each for 100 steps. Compare loss curves and final outputs.

Progress

Progress is local-only — sign in to sync across devices.

← PreviousTrainer and TrainingArguments Next →Evaluation Metrics and the evaluate Library

Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)

No comments yet — be the first.