Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Enhancing Training Efficiency Using Packing with Flash Attention

Achintya Kundu, Rhui Dih Lee, Laura Wynter, Raghu Kiran Ganti, Mayank Mishra

TL;DR

The paper tackles the inefficiency of padding in training variable-length sequences for LLMs by introducing Packing with Position IDs, which concatenates multiple examples and uses position IDs to preserve proper masking. It evaluates online minibatch, offline batch, and optimized sample selection approaches, demonstrating substantial throughput gains—especially with online minibatch packing—while preserving or closely matching loss behavior. The method relies on Flash Attention compatibility and is implemented as PaddingFreeCollator within Hugging Face Transformers 4.44, with broad model support. Across multiple datasets and models, the results show meaningful speedups and nuanced trade-offs between throughput and optimization steps, guiding practical deployment choices for efficient LLM fine-tuning.

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.

Enhancing Training Efficiency Using Packing with Flash Attention

TL;DR

The paper tackles the inefficiency of padding in training variable-length sequences for LLMs by introducing Packing with Position IDs, which concatenates multiple examples and uses position IDs to preserve proper masking. It evaluates online minibatch, offline batch, and optimized sample selection approaches, demonstrating substantial throughput gains—especially with online minibatch packing—while preserving or closely matching loss behavior. The method relies on Flash Attention compatibility and is implemented as PaddingFreeCollator within Hugging Face Transformers 4.44, with broad model support. Across multiple datasets and models, the results show meaningful speedups and nuanced trade-offs between throughput and optimization steps, guiding practical deployment choices for efficient LLM fine-tuning.

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.
Paper Structure (13 sections, 14 equations, 3 figures, 6 tables)

This paper contains 13 sections, 14 equations, 3 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Packing with Position IDs
  4. Online minibatch Collating
  5. Offline Batch Collating
  6. Collating with Optimised Sample Selection
  7. PaddingFreeCollator
  8. Example
  9. Experiments and Results
  10. Main Results
  11. Ablation Study on Minibatch Size
  12. Ablation Study on Sample Selection Method
  13. Conclusions

Figures (3)

  • Figure 1: Histograms of sequence lengths for the 3 training datasets: (top left) FLAN_20k, (top right) OrcaMath_20k, and (bottom) the Stack_20k.
  • Figure 2: Minibatch packing offers lower GPU peak memory usage than padding approach as minibatch size increases. "OOM" denotes out of memory error. Lower is better.
  • Figure 3: Minibatch packing offers better throughput (Tokens/s) than padding as minibatch size increases. "OOM" denotes out of memory error. Higher is better.