使用 Hugging Face Transformers 和 Ray 进行检索增强生成

Hugging Face 已将 Ray 集成到检索增强生成 (RAG) 模型的文档检索机制中。这种集成将检索调用延迟降低了高达 2 倍，并提高了知识密集型 NLP 任务中分布式微调的可扩展性。

Understanding Retrieval Augmented Generation (RAG)

检索增强生成 (RAG) 是一种序列到序列 (seq2seq) 架构，它通过在执行期间从外部知识库（例如 Wikipedia 文本语料库）检索上下文文档来增强其知识。与仅依赖于模型内置参数的标准模型不同，RAG 将这些内部参数与从外部段落检索到的信息相结合，以产生输出。这种双源方法使 RAG 在问答等知识密集型任务中优于其他最先进的模型。

Scaling Distributed Fine-Tuning with Ray

虽然检索步骤对于 RAG 的性能至关重要，但它在分布式微调期间引入了显著的复杂性。在数据并行训练例程中，文档索引通常太大，以至于每个训练工作节点都无法加载副本，从而造成了潜在的瓶颈。

此前，RAG 微调利用 torch.distributed 通信包进行文档检索。然而，这种实现方式有两个主要限制：

  1. Synchronization Bottlenecks: The rank 0 worker was responsible for receiving inputs from all workers, performing the index query, and distributing the results back, which limited performance as the number of training workers increased.
  2. Framework Dependency: The retrieval process group was tied to the training process group, requiring PyTorch to be used for the training process.

通过将 torch.distributed 替换为 Ray——一个用于通用分布式和并行编程的 Python 库——Hugging Face 创建了一个与框架无关的实现。利用 Ray 的有状态 Actor 抽象，多个独立于训练进程的进程可以加载索引并并发处理检索查询，从而消除了 rank 0 瓶颈。

Performance Benchmarks

与 torch.distributed 实现相比，集成 Ray 在多 GPU 微调期间带来了更优越的检索性能。随着 GPU 数量的增加，性能差距进一步扩大，并且增加 Ray 检索进程的数量可以进一步优化速度。

Implementation 2 GPU 3 GPU 4 GPU
torch.distributed 2.12 sec/retrieval 2.62 sec/retrieve 3.438 sec/retrieve
Ray (2 retrieval processes) 1.49 sec/retrieve 1.539 sec/retrieve 2.029 sec/retrieve
Ray (4 retrieval processes) 1.145 sec/retrieve 1.484 sec/retrieve 1.66 sec/retrieve

Note: Benchmarks were conducted over 500 training steps with a per-GPU batch size of 8, measuring the time to retrieve contextual documents on the rank 0 worker.

Implementation and Usage

用户可以使用 Hugging Face 提供的基于 PyTorch Lightning 的微调脚本来实现基于 Ray 的检索。该过程包括安装 ray 和 transformers，并在微调脚本中执行以下配置：

  • Distributed Retriever: Set to ray.
  • Retrieval Workers: Specified via the --num_retrieval_workers flag.

对于寻求进一步优化的用户，可以使用与 Ray Tune 的集成来进行可扩展的超参数调优，以提高模型准确性。

Sources