Causal LM Data Collator

Your tokenizer and dataset pipeline just handed you a perfect batch of examples — but your model is staring at it and refuses to learn. Why? Because for causal language modeling, every batch needs to be shaped into input/label pairs where the label at each position is simply the next token. That’s the job of a data collator. Without it — or with the wrong one — you get misaligned labels, wasted compute on padding tokens, and training loops that silently diverge. In this lesson, you’ll build a data collator from scratch and master the built-in DataCollatorForLanguageModeling so your fine-tunes start fast and stay correct.

The problem this lesson solves

Fine-tuning a causal LM (like GPT‑2, LLaMA, or Mistral) isn’t classification. Your model predicts the next token in a sequence, so every training example is also its own label. The naive approach — grab a batch of tokenized strings and feed them straight to the model — fails miserably:

• The model’s labels argument expects a tensor with the same shape as input_ids, shifted by one.
• If you don’t create that shift, you’re asking the model to predict the current token from itself — guaranteed poor gradients.
• Different sequence lengths in a batch force you to pad, but if you compute loss over padding tokens, you’re teaching the model that <pad> should follow everything.

🩸 The pain is real: a botched collator shows up as a training loss that barely drops, or a model that generates garbage because it learned padding positions.

The data collator is the bridge between your raw dataset and the model’s expected format. It’s where you decide how to pack sequences, how to pad, and what the labels look like — all in one reusable callable.

Core concept / mental model

Think of a data collator as a factory assembly line for training batches. You feed it a list of raw examples (each a dict of token IDs), and it outputs a single batch dictionary with exactly the keys your model’s forward pass expects — input_ids, attention_mask, and labels.

For causal LM, the key transformation is next-token alignment:

• For a sequence [A, B, C, D], the labels are [B, C, D, <pad>] (shifted left by one).
• At training time, position i predicts token i+1.
• The loss is masked so padding tokens don’t contribute.

The built-in DataCollatorForLanguageModeling from transformers does this shift for you, along with optional masking for MLM (not used for causal). But sometimes you need more control — e.g., when you want to pack multiple short sequences into one long context. That’s when you write your own.

Definitions

• data collator: a callable that takes a list of dicts and returns a batch dict of tensors.
• causal LM: a model trained to predict the next token given the left context.
• label shift: moving the target tokens one position right so each position predicts the following token.
• padding: adding a special token so all sequences in a batch have equal length.

How it works step by step

Creating a data collator is a straightforward but critical pipeline, and each step has a clear cause-and-effect on training behavior.

1. Collate raw examples: Start with a list of tokenized dictionaries (from your dataset). Each dict contains input_ids (and often attention_mask).
2. Pad to the longest sequence in the batch (or to a fixed length). Use the tokenizer’s pad_token — if your tokenizer lacks one, set it first.
3. Create the label shift: For each sequence, copy input_ids, then shift it one position to the right — e.g., labels = input_ids[1:] + [pad_token_id].
4. Mask padding in the loss: Set labels to -100 at padding positions. Hugging Face models ignore any token with label -100 in the loss computation.
5. Return a batch dictionary with tensors, plus the attention_mask so the model skips padding when computing attention.

💡 The -100 mask is the universal contract: models like LlamaForCausalLM and GPT2LMHeadModel are designed to ignore it.

Hands-on walkthrough

Let’s build a custom collator for causal LM training with PyTorch. We’ll assume a tokenizer (e.g., from GPT‑2) and a small dataset.

1. Custom collator from scratch

import torch
from transformers import AutoTokenizer

class CausalLMCollator:
    def __init__(self, tokenizer, padding=True, max_length=None):
        self.tokenizer = tokenizer
        self.padding = padding
        self.max_length = max_length

    def __call__(self, features):
        # Handle key names: some datasets use 'input_ids', some use 'text'
        if 'input_ids' not in features[0]:
            features = [self.tokenizer(f["text"], truncation=True, max_length=self.max_length) for f in features]

        # Pad batches
        batch = self.tokenizer.pad(
            features,
            padding=self.padding,
            max_length=self.max_length,
            return_tensors="pt"
        )

        # Create labels shifted by one
        labels = batch["input_ids"].clone()
        if self.tokenizer.pad_token_id is not None:
            labels[labels == self.tokenizer.pad_token_id] = -100

        # Shift labels to the right (next token prediction)
        labels = torch.cat([labels[:, 1:], torch.full((labels.size(0), 1), -100, dtype=torch.long)], dim=1)
        batch["labels"] = labels
        return batch

# Usage
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token  # GPT-2 has no pad token
collator = CausalLMCollator(tokenizer)

sample = [{"input_ids": [1, 2, 3]}, {"input_ids": [4, 5]}]
batch = collator(sample)
print(batch["input_ids"])
print(batch["labels"])

Output (with padding token = 50256, shown as p):

tensor([[1, 2, 3],
        [4, 5, 50256]])
tensor([[2, 3, -100],
        [5, -100, -100]])

2. Using the built-in collator

The quickest route for most fine-tunes:

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,  # crucial: causal LM is not masked LM
    pad_to_multiple_of=8  # efficient padding for TPUs/GPUs
)

# Then pass to Trainer
# trainer = Trainer(..., data_collator=data_collator)

3. Handling raw text inputs

If your dataset stores plain strings, you can tokenize inside the collator:

class TextCollator:
    def __init__(self, tokenizer, max_length):
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __call__(self, features):
        texts = [f["text"] for f in features]
        batch = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt"
        )
        # Shift labels as before
        labels = batch["input_ids"].clone()
        labels[labels == self.tokenizer.pad_token_id] = -100
        labels = torch.cat([labels[:, 1:], torch.full((labels.size(0), 1), -100, dtype=torch.long)], dim=1)
        batch["labels"] = labels
        return batch

📋 Pro tip: Always include attention_mask — it prevents the model from attending to padding and is required for Trainer-based fine-tuning.

Compare options / when to choose what

Troubleshooting & edge cases

Even with a solid collator, things can and will go wrong. Here’s how to diagnose and fix the common failure modes:

• Loss isn’t decreasing — Possibly you didn’t mask padding tokens, so the model learns to predict <pad> everywhere. Fix: set labels to -100 at pad positions.
• Received a Tensor object with TypeError: 'Tensor' object is not callable — You passed a collator that returned a tensor directly instead of a dict. Make sure your __call__ returns a dictionary.
• ValueError: labels shape must match input_ids — If you shift manually, double-check you appended a -100 token; shapes must stay identical.
• tokenizer.pad_token is None — Many tokenizers (e.g., GPT‑2) lack a pad token. Set tokenizer.pad_token = tokenizer.eos_token before calling the collator.
• Padding to the right vs left — For causal LMs, always pad to the right (post-padding). Left-padding breaks the next-token prediction at the start of the sequence.

What you learned & what's next

You now can:

• Explain the role of a data collator in causal LM training.
• Build a custom collator that pads, shifts labels, and masks padding with -100.
• Use the built-in DataCollatorForLanguageModeling for standard cases.
• Choose between built-in, custom, or on-the-fly tokenizing collators.

You’ve also armed yourself with a reliable way to debug misaligned labels and padding issues.

Now that your batches are correct, the next step in this track is to define the training arguments and Trainer — the part that wires your collator, model, and dataset into a full fine-tuning loop. That’s where you’ll set learning rate, batch size, and evaluation strategy.

Keep this collator as your silent partner — it’s the difference between a model that learns your data and one that learns your padding.