Improving Hugging Face Training Efficiency Through Packing with Flash Attention 2

TL;DR

Training with packed instruction tuning examples (without padding) is now compatible with Flash Attention 2 in Hugging Face, thanks to a recent PR and the new DataCollatorWithFlattening

It can provide up to 2x improvement in training throughput while maintaining convergence quality. Read on for the details!

Introduction

Padding input sequences in mini-batches is a usual method to collate inputs during training. However, this introduces inefficiencies because of the irrelevant padding tokens. Packing examples without padding, and using the token position information, is a more efficient alternative. However, previous implementations of packing did not consider example boundaries when using Flash Attention 2, resulting in undesired cross-example attention that reduce quality and convergence.

Hugging Face Transformers now addresses this with a new feature that maintains boundary awareness during packing, alongside the introduction of a new data collator, DataCollatorWithFlattening.

By selecting DataCollatorWithFlattening, Hugging Face Trainer users can now seamlessly concatenate sequences into a single tensor while accounting for sequence boundaries during Flash Attention 2 computations. This is achieved through the flash_attn_varlen_func, which calculates the cumulative sequence lengths in each mini-batch (cu_seqlens).

The same feature is available to Hugging Face SFTTrainer users in the TRL library by setting a new flag, padding_free=True, when calling the data collator DataCollatorForCompletionOnlyLM.

Up to 2x throughput increase

We see significant improvement in training throughput using this feature with the new DataCollatorWithFlattening. The figure below shows the throughput measured in tokens/second during training. In this example, the throughput is the per-GPU average over 8 A100-80 GPU over one epoch of a 20K randomly selected sample from two different instruct tuning datasets, FLAN and OrcaMath.