使用 Hugging Face Transformers and Ray 進行檢索增強生成

Hugging Face 已將 Ray 集成到檢索增強生成 (RAG) 模型的文檔檢索機制中。此集成將檢索調用延遲降低了高達 2 倍，並提高了知識密集型 NLP 任務中分布式細調 (fine-tuning) 的可擴展性。

Understanding Retrieval Augmented Generation (RAG)

檢索增強生成 (RAG) 是一種序列到序列 (seq2seq) 架構，它在執行期間透過從外部知識庫（例如 Wikipedia 文本語料庫）檢索上下文文檔來增強其知識。與僅依賴模型內置參數的標準模型不同，RAG 將這些內部參數與從外部段落檢索到的信息結合起來以產生輸出。這種雙源方法使 RAG 在問答等知識密集型任務中優於其他最先進的模型。

Scaling Distributed Fine-Tuning with Ray

雖然檢索步驟對於 RAG 的性能至關重要，但它在分布式細調期間引入了顯著的複雜性。在數據並行訓練程序中，文檔索引通常太大，以至於每個訓練工作節點 (worker) 都無法加載副本，從而造成了潛在的瓶頸。

先前，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