Dattri-LLM speeds up training data use analysis for large language models
dattri-LLM: A Unified and Efficient Library for Training Data Attribution at LLM Scale
Machine LearningArtificial Intelligence
Summary
Understanding how each piece of training data influences a large AI model's responses is hard and slow. The authors created dattri-LLM, a tool that collects and uses information from the training process efficiently without needing changes to existing setups. It helps people figure out which training examples matter most, works with very big models, and runs much faster than similar tools. This makes it easier to analyze and improve training data for large models.
What this means in practice
- •For machine learning engineers: Run efficient and scalable training data attribution analyses on large language models without modifying training loops.
- •For data platform teams: Implement online data selection methods during model training using reusable gradient operations to speed up dataset curation.
Authors
Shixuan Liu, Tongli Zhou, Junwei Deng, Pingbang Hu, Jiaqi W. Ma
Abstract
Training data attribution (TDA) estimates the contribution of individual training examples to model outputs. Most scalable TDA methods rely on per-example gradients, whose computation and use at LLM scale pose challenges in efficiency, compatibility, and extensibility. We introduce dattri-LLM, a TDA library that makes gradient-based attribution more practical at scale. For efficiency, dattri-LLM uses compact gradient representations and dynamically routes gradient operations based on a cost model. For compatibility, its capture mechanism collects per-example gradients from existing training loops that call backward(), without requiring changes to the loop or its configuration. This includes distributed training with DDP and FSDP and pipelines built with HuggingFace Transformers, TRL, and OLMo. For extensibility, dattri-LLM exposes reusable gradient operations and training-time callbacks for implementing attribution methods and applications. These interfaces support a variety of attribution methods, including gradient similarity, curvature-based influence, and trajectory-based methods, as well as applications that act on gradients during training, such as online data selection. On the same hardware and workload, dattri-LLM achieves 3.2x the throughput of the fastest competing library on average, scales multiple attribution methods to 110B-parameter models across four H200 GPUs, and offers superior attribution fidelity-cost trade-offs across a range of models with different model families and scales.