使用 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 實現文件檢索的主要缺點是它依賴於用於訓練的同一個程序組，並且只有 rank 0 的訓練工作程序將索引載入到記憶體中。
因此，這種實現有一些侷限性
- 同步瓶頸：rank 0 的工作程序必須接收所有工作程序的輸入，執行索引查詢，然後將結果傳送回其他工作程序。這限制了多個訓練工作程序的效能。
- PyTorch 特定：文件檢索程序組必須依賴於用於訓練的現有程序組，這意味著訓練也必須使用 PyTorch。
使用 Ray 實現的文件檢索
為了克服這些限制，我們引入了一種基於 Ray 的新穎的分散式檢索實現。透過 Ray 的有狀態 actor 抽象，使用與訓練程序分離的多個程序來載入索引並處理檢索查詢。有了多個 Ray actor，檢索不再是瓶頸，PyTorch 也不再是 RAG 的必需品。
如下所示，使用基於 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 秒/檢索 |
不同檢索實現的效能比較。對於每個文件檢索實現，我們使用每個 GPU 批處理大小為 8 執行 500 個訓練步驟，並測量 rank 0 訓練工作程序上為每個批次檢索上下文文件所需的時間。結果表明，使用多個檢索程序可以提高效能，尤其是在我們將訓練擴充套件到多個 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！
# Sample script to finetune RAG using Ray for distributed retrieval.
# Add parent directory to python path to access lightning_base.py
export PYTHONPATH="../":"${PYTHONPATH}"
# Start a single-node Ray cluster.
ray start --head
# A sample finetuning run, you need to specify data_dir, output_dir and model_name_or_path
# run ./examples/rag/finetune_rag_ray.sh --help to see all the possible options
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
# Stop the Ray cluster.
ray stop
接下來是什麼？
使用 Huggingface transformers 中的 RAG 和 Ray 檢索實現進行更快的分散式微調，您可以在自己的知識密集型任務上利用 RAG 進行基於檢索的生成。
此外，超引數調整是 transformer 微調的另一個方面，並且可以對準確性產生巨大影響。對於可擴充套件且簡便的超引數調整，請檢視 Ray Tune 庫。透過使用 Ray Tune 與 PyTorch Lightning 的整合，或與 Huggingface transformers 的內建整合，您可以執行實驗來為您的 RAG 模型找到完美的超引數。
最後，請繼續關注 Huggingface 上可能推出的 RAG 的 Tensorflow 實現！
如果您計劃試用 RAG+Ray 整合，請隨時在 Ray Discourse 上分享您的經驗，或加入 Ray 社群 Slack 進行進一步討論——我們很樂意聽到您的聲音！
