论文

dattri-LLM:用于 LLM 规模训练数据归因的统一高效库

dattri-LLM: A Unified and Efficient Library for Training Data Attribution at LLM Scale

AI 基础设施数据基础设施

摘要

训练数据归因 (TDA) 估计各个训练示例对模型输出的贡献。大多数可扩展的 TDA 方法依赖于每个示例的梯度,其在 LLM 规模的计算和使用在效率、兼容性和可扩展性方面提出了挑战。我们引入了 dattri-LLM,这是一个 TDA 库,它使基于梯度的归因在规模上更加实用。为了提高效率,dattri-LLM 使用紧凑的梯度表示,并根据成本模型动态路由梯度操作。为了兼容性,其捕获机制从调用backward()的现有训练循环中收集每个示例的梯度,而不需要更改循环或其配置。这包括使用 DDP 和 FSDP 进行分布式训练,以及使用 HuggingFace Transformers、TRL 和 OLMo 构建的管道。为了实现可扩展性,dattri-LLM 公开了可重用的梯度操作和训练时回调,以实现归因方法和应用程序。这些接口支持各种归因方法,包括梯度相似性、基于曲率的影响和基于轨迹的方法,以及在训练期间作用于梯度的应用程序,例如在线数据选择。在相同的硬件和工作负载上,dattri-LLM 的吞吐量平均是最快竞争库的 3.2 倍,将多种归因方法扩展到跨四个 H200 GPU 的 110B 参数模型,并在具有不同模型系列和规模的一系列模型中提供卓越的归因保真度与成本权衡。 dattri-LLM 的源代码可在 https://github.com/TRAIS-Lab/dattri-llm. 获取