Enhancing Token Filtering Efficiency in Large Language Model Training with Collider

TL;DR

Collider addresses the gap where token filtering could reduce training workload but failed to improve efficiency due to insufficient sparsity and ineffective sparse GEMM. It introduces backward activation filtering across all layers and converts sparse GEMMs into dimension-reduced dense GEMMs via an automatic, graph-updating workflow, implemented as a PyTorch C++ extension with TorchGen support. Empirical results show up to 35.1% backward-time savings and up to 22% end-to-end training reductions when filtering 40% of tokens, along with a 16.3% relative gain in model utility on TinyLlama while reducing training time from 4.5 to 3.7 days on 15B tokens. Collider also promises better scalability to longer contexts and higher filtering ratios, and offers near-seamless integration with existing LLM training pipelines through a single-line change in code.

Abstract

Token filtering has been proposed to enhance utility of large language models (LLMs) by eliminating inconsequential tokens during training. While using fewer tokens should reduce computational workloads, existing studies have not succeeded in achieving higher efficiency. This is primarily due to the insufficient sparsity caused by filtering tokens only in the output layers, as well as inefficient sparse GEMM (General Matrix Multiplication), even when having sufficient sparsity. This paper presents Collider, a system unleashing the full efficiency of token filtering in LLM training. At its core, Collider filters activations of inconsequential tokens across all layers to maintain sparsity. Additionally, it features an automatic workflow that transforms sparse GEMM into dimension-reduced dense GEMM for optimized efficiency. Evaluations on three LLMs-TinyLlama-1.1B, Qwen2.5-1.5B, and Phi1.5-1.4B-demonstrate that Collider reduces backpropagation time by up to 35.1% and end-to-end training time by up to 22.0% when filtering 40% of tokens. Utility assessments of training TinyLlama on 15B tokens indicate that Collider sustains the utility advancements of token filtering by relatively improving model utility by 16.3% comparing to regular training, and reduces training time from 4.7 days to 3.5 days using 8 GPUs. Collider is designed for easy integration into existing LLM training frameworks, allowing systems already using token filtering to accelerate training with just one line of code.

Figures (13)

• Figure 1: An overview of LLM training.
• Figure 2: An overview of existing token filter studies. Forward token filtering methods (a) filter hidden states during forward process, while backward token filtering methods (b) filter training loss during the backward process.
• Figure 3: Leaving the activation (i.e., $softmax$) of filtered tokens unchanged makes the $\mathbf{V}$'s gradients computed by the attention block not sparse anymore after the backpropagation. The dense gradients $\mathbf{G}_V$ will be passed to the front layers, undermining sparsity in all the rest computations .
• Figure 4: PyTorch sparse GEMM outperforms regular GEMM only when filtering more than 95% tokens and cannot improve efficiency of token filtering training which typically drops 30% $\sim$ 40% tokens RHO.
• Figure 5: By filtering the activations in the backward of attention block, we maintain the efficiency advantages of backward token filtering. Other activations (e.g., for backward on $\mathbf{Q}$ and $\mathbf{K}$) are also filtered simultaneously.