Batching Instruction Data: Collate & Loss Masking

This is the chapter’s engineering heart. Formatted instruction texts vary wildly in length, and their padding must not pollute the loss. The solution unfolds in three drafts of a custom collate function — ending with the −100 trick that makes PyTorch’s loss politely ignore padding.

7.3 From dataset items to training batches

InstructionDataset stores each entry as one tokenized text (prompt + response). Unlike ch6’s SpamDataset, it does no padding — that moves to batch-assembly time:

</> InstructionDataset — tokenize the full text, defer the padding

import torch
from torch.utils.data import Dataset

class InstructionDataset(Dataset):
  def __init__(self, data, tokenizer):
      self.data = data

      # Pre-tokenize texts
      self.encoded_texts = []
      for entry in data:
          instruction_plus_input = format_input(entry)
          response_text = f"\n\n### Response:\n{entry['output']}"  #A
          full_text = instruction_plus_input + response_text
          self.encoded_texts.append(
              tokenizer.encode(full_text)
          )

  def __getitem__(self, index):
      return self.encoded_texts[index]  #B

  def __len__(self):
      return len(self.data)

#A Training text = the 7.2 prompt PLUS the ### Response: section with the desired output — the model learns the whole continuation.
#B Items are raw variable-length ID lists. Assembling them into rectangular batches is the collate function's job.

Watch the three collate drafts evolve on the book’s toy batch — three “sequences” of lengths 5, 2 and 3:

The problem: ragged items

InstructionDataset returns raw ID lists of different lengths (5, 2, 3 in the toy). A batch tensor must be rectangular — someone has to pad.

1 / 4

The collate function in three drafts (book pp. 113–119), on the toy batch [0..4], [5,6], [7,8,9]. Every tensor shown is the book's printed output.

</> The final custom_collate_fn (drafts 1 and 2 are subsets)

def custom_collate_fn(
  batch,
  pad_token_id=50256,
  ignore_index=-100,
  allowed_max_length=None,
  device="cpu"
):
  # Find the longest sequence in the batch
  batch_max_length = max(len(item)+1 for item in batch)  #A

  inputs_lst, targets_lst = [], []

  for item in batch:
      new_item = item.copy()
      # Add an <|endoftext|> token
      new_item += [pad_token_id]
      # Pad sequences to max_length
      padded = (
          new_item + [pad_token_id] *
          (batch_max_length - len(new_item))
      )
      inputs = torch.tensor(padded[:-1])   # Truncate the last token for inputs
      targets = torch.tensor(padded[1:])   # Shift +1 to the right for targets  #B

      # Replace all but the first padding tokens in targets by ignore_index
      mask = targets == pad_token_id
      indices = torch.nonzero(mask).squeeze()
      if indices.numel() > 1:
          targets[indices[1:]] = ignore_index  #C

      # Optionally truncate to maximum sequence length
      if allowed_max_length is not None:
          inputs = inputs[:allowed_max_length]
          targets = targets[:allowed_max_length]

      inputs_lst.append(inputs)
      targets_lst.append(targets)

  inputs_tensor = torch.stack(inputs_lst).to(device)  #D
  targets_tensor = torch.stack(targets_lst).to(device)
  return inputs_tensor, targets_tensor

#A The +1 provides room for the shift: after appending one pad and padding to max+1, dropping the last input token and the first target token yields aligned (len max) pairs.
#B The ch2/ch5 next-token setup, applied per batch item.
#C indices[1:] skips the FIRST pad occurrence — that one stays a genuine "response finished" target; the rest become −100.
#D Stack and ship to the device inside the collate — data arrives GPU-ready.

targets after masking (3×5, toy batch)

item 1	1	2	3	4	50256
item 2	6	50256	-100	-100	-100
item 3	8	9	50256	-100	-100

The book's printed final targets. Each row keeps exactly one 50256 (the learnable stop signal); every additional padding position is −100 — invisible to the loss.

The −100 proof

Why does −100 mean “ignore”? The book demonstrates it with a 3-line experiment:

</> Cross-entropy provably skips −100 targets

logits_1 = torch.tensor(
  [[-1.0, 1.0],  # 1st training example
   [-0.5, 1.5]]  # 2nd training example
)
targets_1 = torch.tensor([0, 1])
loss_1 = torch.nn.functional.cross_entropy(logits_1, targets_1)
print(loss_1)   # tensor(1.1269)

logits_2 = torch.tensor(
  [[-1.0, 1.0],
   [-0.5, 1.5],
   [-0.5, 1.5]]  # New 3rd training example
)
targets_2 = torch.tensor([0, 1, 1])
loss_2 = torch.nn.functional.cross_entropy(logits_2, targets_2)
print(loss_2)   # tensor(0.7936) — the 3rd example changed the loss  #A

targets_3 = torch.tensor([0, 1, -100])  #B
loss_3 = torch.nn.functional.cross_entropy(logits_2, targets_3)
print(loss_3)                      # tensor(1.1269)
print("loss_1 == loss_3:", loss_1 == loss_3)  # tensor(True)

#A Adding a third example shifts the average, as expected.
#B Label the third example −100 and the loss returns EXACTLY to the two-example value: PyTorch's cross_entropy defaults to ignore_index=-100, erasing that position from the computation. This is the mechanism the collate function exploits for padding.

Should the instruction tokens be masked too?

This book masks only padding. Some practitioners also mask the instruction tokens so the loss covers just the response — but researchers are divided: Shi et al. (2024, “Instruction Tuning With Loss Over Instructions”) found that not masking the instructions can benefit LLM performance. The book’s simpler choice — train on everything except padding — is also the empirically safer one.

7.4 The data loaders

</> Plugging the collate function into DataLoader

from functools import partial

customized_collate_fn = partial(
  custom_collate_fn,
  device=device,             #A
  allowed_max_length=1024    # the model's context limit
)

from torch.utils.data import DataLoader

num_workers = 0
batch_size = 8
torch.manual_seed(123)

train_dataset = InstructionDataset(train_data, tokenizer)
train_loader = DataLoader(
  train_dataset,
  batch_size=batch_size,
  collate_fn=customized_collate_fn,  #B
  shuffle=True,
  drop_last=True,
  num_workers=num_workers
)
# ... same for validation and test sets

print("Train loader:")
for inputs, targets in train_loader:
  print(inputs.shape, targets.shape)
# torch.Size([8, 61]) torch.Size([8, 61])  #C
# torch.Size([8, 76]) torch.Size([8, 76])
# ...
# torch.Size([8, 66]) torch.Size([8, 66])
# torch.Size([8, 74]) torch.Size([8, 74])
# torch.Size([8, 69]) torch.Size([8, 69])

#A DataLoader calls collate_fn(batch) with ONE argument; partial pre-binds device and allowed_max_length so our five-parameter function fits that interface.
#B The only structural difference from every previous loader in the book: a custom collate.
#C Dynamic padding, visible: 8 examples per batch, but widths 61, 76, 66, 74, 69… — each batch exactly as wide as its longest member.

One real training pair, straight from the loader (abridged — the full prompt “Below is an instruction… Rewrite the sentence using a simile.” plus response):

</> inputs[0] and targets[0] — the toy pattern at real scale

print(inputs[0])
# tensor([21106,  318,  281, 12064,  326, 8477,  257, 4876,  13, 19430,
#         ...
#           464, 5156,  318,  355, 13779,  355,  257, 4936,  13,
#         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256])

print(targets[0])
# tensor([ 318,  281, 12064,  326, 8477,  257, 4876,  13, 19430,  257,
#         ...
#          464, 5156,  318,  355, 13779,  355,  257, 4936,  13, 50256,
#         -100, -100, -100, -100, -100, -100, -100, -100, -100])  #A

#A Exactly the toy structure: targets = inputs shifted by one, ONE trailing 50256 kept as the stop-signal target, and every further padding position masked to −100. (21106 is the token for "Below" — the Alpaca boilerplate's first word.)

Key idea — pad for the tensor, mask for the loss

Padding exists for geometry (rectangular batches); dynamic padding keeps its cost minimal. Masking exists for learning: with ignore_index = −100, the loss sees real text and exactly one end-of-text target per example — never the filler. With data this carefully shaped, training itself is almost anticlimactic.

📝 Check yourself: batching & loss masking

0 / 5

1.How does ch7's padding strategy differ from ch6's — and why?
2.Predict: custom_collate_draft_2 processes the toy item [5, 6] in a batch whose longest member has 5 tokens. What are its input and target rows?
3.What does replacing padding tokens with −100 in the targets accomplish?
4.Why is the FIRST end-of-text token (50256) after each response kept as a real target instead of being masked to −100?
5.Why wrap custom_collate_fn with functools.partial before giving it to the DataLoader?