Packing, document boundaries and long context

The padding problem

If your examples average 300 tokens and your training window is 2,048, then padding every sequence to the window means roughly 85% of the tensor is padding. You pay for those matrix multiplications in full.

Two fixes, with different risk profiles.

Length-grouped batching (group_by_length=True) sorts similar-length examples together so each batch pads to its own longest member rather than to a global maximum. Safe, simple, and it recovers most of the waste.

Packing concatenates examples end to end until the window is full. Zero padding, maximum throughput, and one hazard that is invisible unless you look for it.

Attention across the seam

In a packed sequence, example B follows example A in the same tensor. With ordinary causal attention, the tokens of B can attend to every token of A. The model is being asked to predict the start of one customer’s ticket having read a different customer’s ticket, and it learns a transition that will never occur at inference.

The consequences are subtle rather than catastrophic (slightly worse quality, a model that sometimes continues past the end of an answer) and they never show up in the loss curve.

The correct handling has two parts:

1. A block-diagonal attention mask, so each document attends only within itself. Implemented efficiently through FlashAttention’s variable-length path, which takes cumulative sequence-length offsets rather than a mask tensor.
2. Position ids reset per document, so the second document’s first token is at position 0 rather than position 400. Without this, an example’s position encoding depends on what happened to be packed before it.

Recent TRL versions handle both when packing is enabled, but implementations have differed over time and across libraries. Verify rather than assume:

batch = next(iter(trainer.get_train_dataloader()))
print(batch.keys())                 # look for position_ids or cu_seq_lens
print(batch["position_ids"][0][:50] if "position_ids" in batch else "none")

If position ids climb monotonically across the whole window, documents are not being separated. If you cannot confirm it, use length-grouped batching instead: you lose some throughput and you lose the failure mode entirely.

One cheap mitigation regardless: make sure an end-of-turn or end-of-sequence token sits between packed documents, so there is at least a learned signal that one thing has finished.

Extending a model’s context

A model pretrained at 8,192 tokens does not simply work at 32,768. Its rotary position embeddings encode positions it has never seen, and quality falls off a cliff past the trained length.

Context extension fine-tuning changes the position encoding and then trains briefly at the longer length. The usual approaches all adjust RoPE’s frequency base:

config.rope_scaling = {"type": "yarn", "factor": 4.0,
                       "original_max_position_embeddings": 8192}

Linear interpolation is the simplest and loses the most local resolution; NTK-aware scaling and YaRN adjust different frequency bands differently and preserve short-range behaviour better. All of them want a short fine-tune on long sequences afterwards: typically a few hundred steps, which is cheap in steps and expensive in memory, because attention is quadratic and you are now at four times the length.

The data problem nobody solves well

Here is the honest difficulty: long-context supervised data barely exists. Real conversations are short. Real tickets are short. If 95% of your examples are under 1,000 tokens and you train with a 32,768 window, you have paid for a long context and taught the model nothing about using one.

The common synthetic route is to build multi-document tasks: concatenate several related documents and ask a question whose answer requires one of them, or several. It works, and it has a well-documented failure. If all your long-context training is of the “find the fact in the haystack” kind, the model becomes good at retrieval within its context and no better at synthesising across it: summarising a long document, tracking an entity through a transcript, reconciling two sources that disagree.

So build both kinds, and evaluate both kinds. And before any of it, check the sequence-length histogram of your data against the window you are about to pay for. If the histogram says you do not need the window, do not buy it.