You turned on sequence packing, throughput went up 3x, the loss curve got visibly smoother, and your eval got worse. Not catastrophically worse — just enough that you blamed the learning rate, then the data mix, then the seed.
The learning rate is fine. Your packed batches let token 4,000 attend to a completely unrelated training example that happens to sit earlier in the same 8,192-token buffer. The causal mask permits it, FlashAttention happily computes it, and the model learns a conditional distribution that will never exist at inference time.
TL;DR
- Sequence packing concatenates short training samples into one long buffer to kill padding waste. Done naively, the causal mask lets each sample attend to every earlier sample in the buffer — cross-document attention.
- Cross-document attention lowers training loss (extra context is free signal) while raising eval loss, because at inference there is no foreign prefix. Your loss curve lies to you.
- With 300-token median samples packed into 8,192, roughly 96% of your tokens are trained with a foreign prefix. This is not an edge case.
- The fix is block-diagonal attention:
cu_seqlenswithflash_attn_varlen_func, or aBlockMaskwith FlexAttention, plus per-documentposition_idsreset. - Verify it numerically. Run one sample alone, run it packed, and diff the logits. If they don't match to ~1e-3, your mask is not doing what you think.
Why does everyone use sequence packing in the first place?
Because padding is pure waste and real SFT datasets are viciously skewed. A typical instruction-tuning set has a median around 200–400 tokens and a tail out to 8k+. If you pad every sequence to the max length in the batch, you burn compute on padding tokens for the whole MLP stack (linear in token count) and, unless you use a varlen kernel, for attention too.
The waste ratio is just 1 - mean_len / max_len. With a median of 300 and a bucket max of 4,096, you're at over 90% padding on most batches. Packing recovers essentially all of it. That is the whole appeal, and it's real — this is why every serious training stack does it.
The problem is what packing does to the attention mask, not to the throughput.
What exactly does cross-document attention corrupt?
Take a packed buffer holding samples A, B, C in that order. A standard causal mask says "token i may attend to all j ≤ i." So every token in B can see all of A. Every token in C sees A and B.
The model is now trained on P(token | its own prefix + an arbitrary unrelated document). At inference you serve P(token | its own prefix). That's a train/serve mismatch on the conditioning set, which is the one thing a language model is entirely made of.
Two concrete ways this bites:
1. Format leakage. SFT samples share structure — the same system prompt, the same ### Response: marker, the same JSON envelope. When B can see A, the model can copy the completed structure from A instead of learning to generate it from its own instruction. Copying is a much easier gradient path than induction. You are training a retrieval head where you wanted a generation policy.
2. Answer leakage across near-duplicates. Datasets built by templating or by augmenting a seed set contain near-duplicates. Shuffling does not prevent two paraphrases of the same QA pair from landing in one buffer. When it happens, the loss on the second one drops to near zero and your training curve records that as learning.
Here's the part that catches people: cross-document attention can only help training loss. The extra context is either useful or ignorable; softmax attention can always learn to down-weight it. So the broken configuration produces a strictly nicer-looking loss curve than the correct one. The curve is not evidence.
How many of my tokens are actually affected?
Compute it, don't guess. If your median sample is m tokens and your pack length is L, each buffer holds about k = L/m documents, and every token outside the first document sees foreign context:
affected ≈ 1 - 1/k = 1 - m/L
- m=300, L=8192 → k≈27 → 96% of tokens
- m=2000, L=8192 → k≈4 → 75% of tokens
- m=4000, L=8192 → k≈2 → 50% of tokens
There is no packing configuration where this is a small correction. The only regime where it's near-zero is when your samples are already close to the pack length, which is exactly the regime where packing buys you nothing.
How do I implement block-diagonal attention correctly?
You need two things: a varlen attention path and per-document position_ids. Here is the collator side — build the flat buffer and the cumulative-sequence-length index that FlashAttention wants:
import torch
def pack(samples, max_len=8192):
"""samples: list of dicts with 'input_ids' and 'labels' (‑100 = masked)."""
ids, labels, pos, cu = [], [], [], [0]
for s in samples:
n = len(s["input_ids"])
if len(ids) + n > max_len:
break
ids.extend(s["input_ids"])
labels.extend(s["labels"])
pos.extend(range(n)) # reset per document, NOT 0..max_len
cu.append(len(ids))
return {
"input_ids": torch.tensor(ids, dtype=torch.long),
"labels": torch.tensor(labels, dtype=torch.long),
"position_ids":torch.tensor(pos, dtype=torch.long),
"cu_seqlens": torch.tensor(cu, dtype=torch.int32),
"max_seqlen": max(cu[i+1]-cu[i] for i in range(len(cu)-1)),
}
And the attention call. flash_attn_varlen_func takes unpadded (total_tokens, n_heads, head_dim) tensors and treats each cu_seqlens segment as an independent sequence — that is the block-diagonal mask, with no N×N matrix ever materialized:
from flash_attn import flash_attn_varlen_func
out = flash_attn_varlen_func(
q, k, v, # (total_tokens, n_heads, head_dim)
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
causal=True,
)
If you'd rather stay in pure PyTorch, FlexAttention expresses the same thing declaratively and compiles to a sparse kernel:
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
doc_id = torch.repeat_interleave(
torch.arange(len(lengths), device="cuda"),
torch.tensor(lengths, device="cuda"),
)
def doc_causal(b, h, q_idx, kv_idx):
return (doc_id[q_idx] == doc_id[kv_idx]) & (q_idx >= kv_idx)
block_mask = create_block_mask(doc_causal, B=None, H=None,
Q_LEN=total, KV_LEN=total)
out = flex_attention(q, k, v, block_mask=block_mask)
What you should not do is build a dense 4D boolean mask and hand it to scaled_dot_product_attention. At L=8192 that mask alone is ~67M elements per sequence, and you've just given back the memory savings that packing earned.
Do I really need to reset position_ids with RoPE?
Partially — and this is where a lot of blog advice is wrong. RoPE encodes relative offsets, so if document B occupies absolute positions 3000–3300 contiguously, the relative distances inside B are already correct. Resetting positions to 0–300 does not fix a broken attention mask, and a correct attention mask already gives you the right relative geometry.
Reset anyway, for three reasons:
-
It's the boundary signal. In Hugging Face's padding-free path, resetting
position_idsis how the varlencu_seqlensgets derived. Non-reset positions there mean no block-diagonal mask at all. - Absolute position does matter to a trained model. Attention sinks form at low absolute positions, and RoPE's long-range decay means a model trained where every document starts at a different large offset behaves differently at position 0, which is exactly where every inference request starts.
- Any absolute or learned positional component (ALiBi slopes aside, some architectures interleave NoPE layers) breaks outright without the reset.
Why did my loss get bigger after I fixed the mask?
It should. You removed free context. A higher training loss with a correct mask is the expected outcome and usually comes with better held-out numbers. Judge the fix on eval, never on the train curve.
While you're in there, check one more thing that packing silently changes: loss normalization. Unpacked, most trainers average per-sequence then average across the batch, so every sample carries weight 1/B. Packed, you get one flat token-mean, so a 2,000-token sample carries 20x the weight of a 100-token one. That's a different objective. It also interacts badly with gradient accumulation — averaging per-microbatch means and calling it the full-batch mean is only correct when every microbatch has the same token count, which packing guarantees it won't. Sum the loss and divide by the global token count, or explicitly re-weight per document if you want uniform per-sample weighting.
How do I verify the mask is actually block-diagonal?
Numerically, in about ten lines. Run one document standalone, run it as the second element of a pack, and compare logits:
solo = model(input_ids=b_ids[None], position_ids=torch.arange(len(b_ids))[None])
packed = model(input_ids=torch.cat([a_ids, b_ids])[None],
position_ids=torch.cat([torch.arange(len(a_ids)),
torch.arange(len(b_ids))])[None],
**varlen_kwargs)
delta = (packed.logits[0, len(a_ids):] - solo.logits[0]).abs().max()
print(delta) # correct: ~1e-3 in bf16. broken: 1.0+
A broken mask does not produce a subtly larger delta. It produces an obviously large one. This test takes five minutes and is the only thing standing between you and a week of blaming your learning rate.
When is naive packing actually fine?
Pretraining. On a continuous web corpus at trillion-token scale, documents in a buffer are unrelated, EOS is a reliable reset signal the model learns to exploit, and every position is seen with countless different prefixes, so the spurious conditioning averages out. Plenty of strong base models were pretrained with plain causal packing.
SFT is the opposite regime: small data, high structural redundancy, few epochs, and samples that share templates. There is nothing to average out over. The same is true for preference tuning and for any distillation set generated from one prompt template.
So why does cross-document attention ruin your SFT?
Because sequence packing without a block-diagonal mask trains the model on a conditioning set that inference will never reproduce: each sample sees an arbitrary earlier sample as its prefix. That extra context can only lower training loss, so the curve improves while the model quietly learns to copy structure from neighbors instead of generating it — and with a 300-token median packed into 8,192, this touches roughly 96% of your tokens. Fix it with cu_seqlens varlen attention or a FlexAttention BlockMask, reset position_ids per document, check your loss normalization now that samples have unequal weight, and confirm the whole thing with a solo-vs-packed logit diff. Expect training loss to rise and eval to improve.
Top comments (0)