Learn AI Create the Future — with Pranay
Main site

The HuggingFace Stack

Data collators

A collator is the function that turns a handful of individual examples into one padded tensor batch, and picking the right one decides what your labels look like.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

A collator gathers a few prepared examples and packs them into one rectangle of numbers for the model.

Think of a tiffin service packing lunch boxes. Each order is a different amount of food, but the delivery crate holds fixed-size boxes. The packer fills each box, tops up the gaps, and hands over one neat crate. Nobody redesigns the crate for every order.

Your examples are the orders. The batch is the crate. The collator is the packer, and it works one crate at a time.

Why it exists

Sentences have different lengths. A batch must be rectangular — the problem from padding and attention masks. Somebody has to do the filling.

Doing it up front, during preparation, means padding everything to one global length. Every short sentence then carries blanks through every epoch, wasting real compute.

Doing it at the last moment is better. Each batch is padded only to its own longest member. A batch of short reviews stays short. This is dynamic padding, and the collator is where it happens.

There is a second job. Some training setups need labels invented on the spot — hiding random words for the model to guess, for instance. A collator can do that too, freshly, every epoch.

How it works

three prepared examples          one batch
  6 tokens  ─┐
 13 tokens  ─┼─→ collator →  a 3 × 13 rectangle, blanks marked in the mask
  3 tokens  ─┘

drop the long one:
  6 tokens  ─┐
  3 tokens  ─┴─→ collator →  a 2 × 6 rectangle — smaller, cheaper

The same collator gives you a different rectangle each time, sized to whatever landed in that batch.

A real example you have seen

Courier packaging. The shop picks a box that fits your order and stuffs paper into the gaps. They do not ship every order in the largest box available — that would cost more for nothing.

Remember this

  • A collator turns a list of examples into one padded batch of tensors.
  • Dynamic padding sizes each batch to its own longest example, saving compute.
  • Some collators also create labels — masking words, or shifting text for next-word training.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install transformers torch

Tested with transformers 5.6 and torch 2.5. Only the tokenizer files of google/bert_uncased_L-2_H-128_A-2 are needed, about 18 MB. All CPU.

Dynamic padding, watched happening

collate.py
from transformers import AutoTokenizer, DataCollatorWithPadding

tok = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
collate = DataCollatorWithPadding(tok)

texts = ["cold biryani",
         "the filter coffee here is genuinely the best in the city",
         "slow"]
rows = [dict(tok(t, truncation=True, max_length=32), labels=lab)
        for t, lab in zip(texts, [0, 1, 0])]
for i, r in enumerate(rows):
    print(f"row {i}: {len(r['input_ids'])} tokens")

batch = collate(rows)
print("batch keys :", sorted(batch.keys()))
print("batch shape:", tuple(batch["input_ids"].shape))
print("attention_mask:\n", batch["attention_mask"])
print("labels:", batch["labels"])

short = collate([rows[0], rows[2]])
print("without the long row:", tuple(short["input_ids"].shape))
Output
row 0: 6 tokens
row 1: 13 tokens
row 2: 3 tokens
batch keys : ['attention_mask', 'input_ids', 'labels', 'token_type_ids']
batch shape: (3, 13)
attention_mask:
 tensor([[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
labels: tensor([0, 1, 0])
without the long row: (2, 6)

The identical collator produced a 3 × 13 batch and a 2 × 6 batch. That second rectangle is under a third the size of the first, for the same two rows.

Collators that build labels for you

label_collators.py
import torch
from transformers import (AutoTokenizer, DataCollatorForLanguageModeling,
                          DataCollatorForTokenClassification)

torch.manual_seed(0)
tok = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")

mlm = DataCollatorForLanguageModeling(tok, mlm=True, mlm_probability=0.3)
b = mlm([tok("the dosa was crisp and the sambar was hot enough"), tok("chai please")])
print("mask token id:", tok.mask_token_id)
print("input_ids[0]:", b["input_ids"][0].tolist())
print("labels[0]   :", b["labels"][0].tolist())
print("labels[1]   :", b["labels"][1].tolist())

ner = DataCollatorForTokenClassification(tok)
f1 = dict(tok("mumbai is big"), labels=[-100, 3, 0, 0, -100])
f2 = dict(tok("delhi"),         labels=[-100, 3, -100])
out = ner([f1, f2])
print("ids   :", out["input_ids"].tolist())
print("labels:", out["labels"].tolist())
Output
mask token id: 103
input_ids[0]: [101, 1996, 103, 21937, 2001, 15594, 1998, 1996, 29086, 2099, 2001, 2980, 103, 102]
labels[0]   : [-100, -100, 9998, 2050, -100, -100, -100, -100, -100, -100, -100, -100, 2438, -100]
labels[1]   : [-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100]
ids   : [[101, 8955, 2003, 2502, 102], [101, 6768, 102, 0, 0]]
labels: [[-100, 3, 0, 0, -100], [-100, 3, -100, -100, -100]]

Masking is random, so your exact positions will differ even with the same seed across library versions. The pattern is what reproduces.

The walkthrough

Read position 2 and position 3 of the masked row together. Position 2 became 103, the [MASK] id, and its label 9998 records the word that was hidden. Position 3 was not masked — it holds 21937 — yet its label is 2050. That row was corrupted with a random word instead of a mask. The classic BERT recipe replaces 80% of chosen tokens with [MASK], 10% with a random token and leaves 10% alone, so the model cannot learn "only predict where you see a mask".

-100 means "do not score this position". PyTorch's cross-entropy loss ignores it by default. Every unchosen token, and every pad, is -100. Row 1 came back all -100: with two real tokens and a 30% chance each, nothing was picked. Short rows sometimes contribute no loss at all, and that is correct behaviour rather than a bug.

Labels get padded with -100; inputs get padded with 0. Look at the last two lines. DataCollatorForTokenClassification pads input_ids with the pad id and labels with the ignore value. Padding labels with 0 instead would silently teach the model that every blank is class 0 — a real bug that produces a model quietly biased toward its first label.

mlm=False switches the same class to causal training. It then copies input_ids into labels and lets the model do the shifting internally. Use it for next-word training as in how LLMs work, and for the fine-tuning loop in QLoRA.

pad_to_multiple_of=8 is a free speed knob on GPU. Tensor cores prefer dimensions that are multiples of 8 (16 for fp16 on some hardware). Rounding the sequence length up costs a few pad tokens and can measurably speed up mixed-precision training.

DataCollatorForSeq2Seq is the fourth one you will meet. Translation and summarisation have inputs and targets of different lengths, so it pads them separately and can prepare decoder inputs. Reach for it whenever the model is encoder-decoder.

Common mistakes

Padding in map, then adding a collator too. The rows are already the same length, so dynamic padding has nothing left to do and you pay full cost forever. Tokenize without padding in map; let the collator pad.

Leaving string columns in the dataset. A collator tries to stack every key it is given. A leftover text column raises an unhelpful error about converting strings to tensors. remove_columns at map time is the fix.

Calling the label column something else. These collators, and the model's forward signature, look for labels. A column named label is renamed for you by the Trainer; target or y is not, and the model then computes no loss at all — training "runs" and learns nothing.

Using DataCollatorForLanguageModeling(mlm=True) with a decoder model. GPT-style models have no [MASK] token, and the collator fails or produces nonsense. Masked training is for encoders like BERT; decoders take mlm=False.

Ignoring batch composition. Dynamic padding only pays off when batchmates are similar lengths. One 512-token outlier stretches its whole batch. group_by_length=True in the Trainer buckets similar lengths together and often cuts training time noticeably.

Try it yourself

Build ten sentences of wildly different lengths, batch them two ways — sorted by length, and shuffled — and add up the total padded cells in each. How much compute did sorting save? Then set pad_to_multiple_of=8 and see what the totals become.

What to learn next

Researcher — Mathematics and papers.

The collator's place in the pipeline

Trainer and PyTorch's DataLoader accept a collate_fn: List[Example] -> Dict[str, Tensor]. It runs in the worker process, after Dataset.__getitem__ and before the batch is pinned and moved to the device — so its cost overlaps compute when num_workers > 0, and lands on the critical path when it is 0. See DataLoader internals and custom collate_fn for the framework-level view; the HuggingFace classes are ordinary collate_fns with tokenizer-aware defaults.

The division of labour is a purity argument. Deterministic per-example work is cached once by map. Batch-dependent work (padding) and intentionally stochastic work (MLM masking, augmentation) must happen per batch, per epoch, and so live here.

The masking objective, precisely

BERT's masked language modelling (Devlin et al., 2019) selects a fraction p of token positions — 15% by default, 30% in the demo to make the effect visible — and for each selected position applies: [MASK] with probability 0.8, a uniformly random vocabulary token with probability 0.1, and the original token with probability 0.1. Loss is computed only over selected positions; everything else is -100.

The 80/10/10 split exists because [MASK] never appears at fine-tuning or inference time. Without the random and unchanged cases, the encoder can learn a representation valid only where a mask token is present — a train/serve mismatch. Later work sharpened the recipe: RoBERTa (Liu et al., 2019) resamples the mask each epoch rather than fixing it once, which is precisely what doing this in the collator gives you for free; SpanBERT (Joshi et al., 2020) masks contiguous spans; ELECTRA (Clark et al., 2020) replaces masking with replaced-token detection, gaining sample efficiency by scoring every position instead of 15% of them.

Padding waste, quantified

For a batch of lengths L₁..L_B, the padded cost is B·L_max while useful work is ΣLᵢ, so utilisation is ΣLᵢ/(B·L_max). Under a heavy-tailed length distribution — the norm for real text — L_max is dominated by the batch maximum, and utilisation falls as B grows. Length-grouped sampling reduces E[L_max] per batch at the cost of correlated batches, which slightly biases gradient estimates; in practice the trade is accepted, and the sampler shuffles within buckets to soften it.

The stronger fixes remove the rectangle entirely: sequence packing with reset causal masks, and varlen attention kernels that take concatenated sequences plus cumulative-length offsets (FlashAttention, Dao et al., 2022). At that point the collator's job shrinks to building the offset index — a different function, same position in the pipeline.

What to learn next