Flash Attention 2 Packing Support Added

A new data collator utilizing Packing technology has been introduced to the Hugging Face Transformers and TRL libraries.

The existing packing method had a problem where training quality degraded when using Flash Attention 2 because it could not distinguish boundaries between examples, but this update enables boundary awareness.

Key Results:

• Training Throughput: On the FLAN dataset, the Llama 2-7B and Mistral-7B models showed up to a 2x increase in throughput.
• Memory Usage: Peak memory usage was reduced by about 20% on the FLAN dataset.
• Training Quality: Compared to the existing method using Padding, there is no difference in Validation Loss, maintaining the same convergence quality.

Usage:

• transformers library: Use DataCollatorWithFlattening.
• trl library: Set the padding_free=True flag when calling DataCollatorForCompletionOnlyLM.