sdadas / unpad-impl

Unpadding implementation for transformers

This repository contains the implementation of automatic unpadding and sequence packing for Hugging Face Transformers models. In version 5 of Transformers, attention handling across models was standardized, enabling older architectures such as BERT and RoBERTa to use Flash Attention. With Flash Attention support, it becomes possible to implement the unpadding trick previously utilized in ModernBERT (see Unpadding and Sequence Packing ). This technique involves automatically removing padding tokens and transforming the input tokens, typically represented as an N x M matrix (where N is the batch size and M is the length of the longest sequence in the batch), into a single, long vector containing the concatenated, padding-free sequences. This vector is then processed by the embedding layer and attention blocks. Right before the final model layer (e.g., LM head, classification head), it is transformed back into a matrix format.

Unpadding can yield significant benefits in terms of computational efficiency and memory savings. The most substantial gains are achieved when there is high variance in the lengths of the processed texts. If the texts are of roughly equal length, the benefits will be minimal, though the model should still run faster than the original implementation. Below are several benchmarks supporting these claims.

Inference

The experiment involved performing reranking using the polish-reranker-roberta-v3 model on a subset of queries from the PIRB benchmark . We used 9 datasets from the Web Datasets group, with a maximum of 1,000 queries per dataset. Each query had 100 candidate documents. The model's task was to evaluate the relevance of the query-document pairs, resulting in over 800,000 predictions in total. The test was conducted on a single NVIDIA RTX A6000 GPU. Across all tested implementations, we used a fixed batch size of 32. The results are presented in the table below.

Fine-Tuning

In the next experiment, we fine-tuned the polish-roberta-8k model on classification tasks. We measured the total time required to train the model for 10 epochs, including the time taken for evaluation on the validation split after each epoch and evaluation on the test set upon training completion. Two datasets were selected for this experiment: POLEMO-IN from the KLEJ benchmark (short to medium text lengths) and banking-long from the FinBench benchmark (medium to very long texts). The tests were conducted on a single NVIDIA RTX A6000 GPU. The results are presented in the table below (OOM - CUDA out-of-memory error). For the dataset with shorter texts, training with unpadding is over twice as fast and consumes three times less memory compared to the default SDPA implementation. For the dataset with long texts, unpadding allows for training with 16 times larger batch sizes and is over three times faster than the original model.