Enhancing Training Efficiency Using Packing with Flash Attention

TL;DR

This paper discusses an improved method for training large language models by using packing with flash attention, which reduces padding inefficiencies and enhances GPU resource utilization.

Contribution

It introduces a new masking feature for packing in Hugging Face Transformers 4.44, enabling more efficient training of LLMs.

Findings

Improved GPU utilization with packing and masking.

Reduced training time due to efficient padding handling.

Enhanced training efficiency across various packing configurations.

Abstract

Padding is often used in tuning LLM models by adding special tokens to shorter training examples to match the length of the longest sequence in each batch. While this ensures uniformity for batch processing, it introduces inefficiencies by including irrelevant padding tokens in the computation and wastes GPU resources. Hugging Face SFT trainer has always offered the option to use packing to combine multiple training examples, allowing for maximal utilization of GPU resources. However, up till now, it did not offer proper masking of each packed training example. This capability has been added to Hugging Face Transformers 4.44. We analyse this new feature and show the benefits across different variations of packing.