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

来自Anyscale团队的客座博客文章，作者为Amog Kamsetty

Huggingface Transformers最近增加了检索增强生成（RAG）模型，这是一种新的自然语言处理架构，它利用外部文档（如维基百科）来增强其知识，并在知识密集型任务上实现了最新的成果。在这篇博客文章中，我们介绍了将用于构建可扩展应用程序的Ray库集成到RAG上下文文档检索机制中。这将检索调用速度提高了2倍，并改善了RAG分布式微调的可扩展性。

什么是检索增强生成（RAG）？

RAG的概述。该模型从外部数据集中检索上下文文档作为其执行的一部分。这些上下文文档与原始输入一起用于生成输出。该GIF来源于Facebook的原始博客文章。

最近，Huggingface与Facebook AI合作，在其Transformers库中引入了RAG模型。

RAG表现得就像任何其他的seq2seq模型一样。然而，RAG有一个中间组件，它从外部知识库（如维基百科的文本语料库）中检索上下文文档。然后，这些文档与输入序列一起使用，并传递给底层的seq2seq生成器。

这个信息检索步骤使RAG能够利用多个知识源，包括嵌入模型参数和包含在上下文段落中的信息，使其在问答等任务中胜过其他最新的模型。您可以使用Huggingface提供的这个演示来亲自尝试。

扩展微调

这种上下文文档的检索对于RAG的最新成果至关重要，但也引入了额外的复杂性。当通过数据并行训练过程扩展训练时，一个简单的文档查找实现可能成为训练的瓶颈。此外，检索组件中使用的文档索引通常相当大，使每个训练工作者加载自己的复制索引变得不可行。

之前的RAG微调实现利用了torch.distributed通信包进行文档检索部分。然而，这个实现有时被证明是不灵活且有限的可扩展性。

相反，需要一个与框架无关且更灵活的特殊编程实现。Ray正好符合这个要求。Ray是一个简单而强大的Python库，用于通用分布式和并行编程。使用Ray进行分布式文档检索，我们实现了与torch.distributed相比每次检索调用的2倍速度提升，并且整体上提升了微调的可扩展性。

Ray用于文档检索

使用torch.distributed实现的文档检索

torch.distributed实现的文档检索的主要缺点是它依赖于用于训练的同一进程组，只有0号训练工作者将索引加载到内存中。

因此，这个实现有一些限制：

同步瓶颈：0号工作者必须接收来自所有工作者的输入，执行索引查询，然后将结果发送回其他工作者。这限制了多个训练工作者的性能。
PyTorch特定：文档检索过程组必须依赖于用于训练的现有进程组，这意味着训练也必须使用PyTorch。

使用Ray实现的文档检索

为了克服这些限制，我们引入了一种基于Ray的新型分布式检索实现。使用多个Ray actor，这些actor与训练进程分开，用于加载索引和处理检索查询。使用多个Ray actor，检索不再是瓶颈，而且RAG不再需要PyTorch。

正如您可以在下面看到的，使用基于Ray的实现可以提高多GPU微调的检索性能。下面的结果显示每次检索调用的秒数，可以看出当我们增加训练的GPU数量时，使用Ray的性能相对于torch.distributed更好。另外，如果我们增加执行检索的Ray进程数，由于单个检索进程不再是瓶颈，使用更多的训练工作者也会获得更好的性能。

	2 GPU	3 GPU	4 GPU
torch.distributed	2.12 秒/检索	2.62 秒/检索	3.438 秒/检索
Ray 2 检索进程	1.49 秒/检索	1.539 秒/检索	2.029 秒/检索
Ray 4 检索进程	1.145 秒/检索	1.484 秒/检索	1.66 秒/检索

不同检索实现的性能比较。对于每个文档检索实现，我们在排名0的训练节点上运行500个训练步骤，每个GPU批处理大小为8，并测量检索每个批次的上下文文档所需的时间。结果显示，使用多个检索进程可以提高性能，特别是在将训练扩展到多个GPU时。

如何使用？

Huggingface提供了一个基于PyTorch Lightning的微调脚本，并我们扩展了它，以将Ray检索实现作为一个选项添加进去。

要试用它，请首先安装必要的依赖：

pip install ray
pip install transformers
pip install -r transformers/examples/research_projects/rag/requirements.txt

然后，您可以指定您的数据路径和其他配置，并运行finetune-rag-ray.sh！

# 通过Ray进行分布式检索的RAG微调示例脚本。

# 将父目录添加到Python路径以访问lightning_base.py
export PYTHONPATH="../":"${PYTHONPATH}"

# 启动单节点Ray集群。
ray start --head

# 一个示例微调运行，您需要指定data_dir、output_dir和model_name_or_path
# 运行./examples/rag/finetune_rag_ray.sh --help以查看所有可能的选项

python examples/rag/finetune_rag.py \
    --data_dir $DATA_DIR \
    --output_dir $OUTPUT_DIR \
    --model_name_or_path $MODEL_NAME_OR_PATH \
    --model_type rag_sequence \
    --fp16 \
    --gpus 8 \
    --profile \
    --do_train \
    --do_predict \
    --n_val -1 \
    --train_batch_size 8 \
    --eval_batch_size 1 \
    --max_source_length 128 \
    --max_target_length 25 \
    --val_max_target_length 25 \
    --test_max_target_length 25 \
    --label_smoothing 0.1 \
    --dropout 0.1 \
    --attention_dropout 0.1 \
    --weight_decay 0.001 \
    --adam_epsilon 1e-08 \
    --max_grad_norm 0.1 \
    --lr_scheduler polynomial \
    --learning_rate 3e-05 \
    --num_train_epochs 100 \
    --warmup_steps 500 \
    --gradient_accumulation_steps 1 \
    --distributed_retriever ray \
    --num_retrieval_workers 4

# 停止Ray集群。
ray stop

下一步是什么？

使用Huggingface transformers和Ray检索实现的RAG进行更快的分布式微调，您可以在自己的知识密集型任务上利用RAG进行基于检索的生成。

此外，超参数调整是transformer微调的另一个方面，对准确性有巨大影响。要进行可伸缩且易于调整的超参数调整，请查看Ray Tune库。通过使用Ray Tune与PyTorch Lightning的集成，或与Huggingface transformers的内置集成，您可以运行实验来找到适合您的RAG模型的完美超参数。

最后，敬请期待Huggingface上RAG的潜在Tensorflow实现！

如果您打算尝试RAG+Ray集成，请随时在Ray Discourse上分享您的经验，或加入Ray社区的Slack进行进一步讨论–我们很乐意听取您的意见！

也发布在 https://medium.com/distributed-computing-with-ray/retrieval-augmented-generation-with-huggingface-transformers-and-ray-b09b56161b1e

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