Hugging Face packs 2x training speed bump into Flash Attention 2 with a single collator swap

Padding tokens have always been the silent performance killer in LLM training. You batch your examples, stuff them with empty tokens to match the longest sequence, and then watch GPUs waste cycles computing attention on nothing. Hugging Face just rolled out a fix that feels overdue: the new DataCollatorWithFlattening, which ditches padding entirely by packing examples into a single concatenated tensor.

This isn’t just another token-dropping trick. The collator works with Flash Attention 2’s flash_attn_varlen_func, which tracks cumulative sequence lengths so the model never attends across example boundaries. That boundary-awareness was the missing piece—previous packing attempts let attention bleed between unrelated samples and tanked convergence. Now, the packed version produces identical validation loss curves as traditional padded training on Llama 2 7B, Mistral 7B, and Granite 8B Code across both FLAN and OrcaMath datasets.

The numbers are substantial. On FLAN—a dataset with short but wildly variable-length sequences—throughput doubled on 8 A100-80 GPUs while peak memory dropped 20%. On the more uniform OrcaMath dataset, gains were a still-respectable 1.4x throughput increase and 6% memory savings. You get these wins without changing your optimization steps, since the collator preserves the original mini-batch structure rather than cramming more examples into fewer steps.

Fourteen model architectures already expose the necessary position_ids, including the Llama family, Mistral, Mixtral, Gemma, Phi, and Qwen 2. If you’re using Hugging Face Trainer, it’s a two-line change: instantiate your model with Flash Attention 2 and swap in DataCollatorWithFlattening. TRL users working with DataCollatorForCompletionOnlyLM just flip padding_free=True. The barrier to adoption is practically zero, which suggests this will become the default way sensible teams train instruction-tuned models going forward.