# Retrieval Augmented Generation with Huggingface Transformers and Ray

huggingface.co · system · Feb 10, 2021

**Type:** Problem & solution

## Summary

Huggingface Transformers added the Retrieval Augmented Generation architecture to combine standard sequence-to-sequence models with external document retrieval for knowledge-intensive tasks. Scaling RAG fine-tuning across multiple GPUs previously created bottlenecks because the torch.distributed implementation forced the single rank 0 worker to handle all retrieval index lookups. To resolve this synchronization bottleneck, Ray was integrated into the contextual document retrieval mechanism using stateful actor abstractions. Dedicated Ray processes separate from the training workers now load the large index and process retrieval requests concurrently. This architectural separation delivers up to a 2x speedup per retrieval call and removes the strict dependency on PyTorch for training.

## Context

When scaling RAG fine-tuning via data-parallel training routines, naive document retrieval lookups become a bottleneck. The document index is too large for every training worker to load its own copy into memory, and the previous torch.distributed implementation latched onto the training process group, forcing the rank 0 worker to synchronously receive all worker inputs, query the index, and return results.

## Approach / What changed

Ray was integrated into RAG's document retrieval mechanism using stateful actor abstractions. Separate Ray actor processes load the document index and handle retrieval queries independently from training workers, avoiding the single-worker bottleneck and decoupling retrieval from PyTorch.

## Takeaways

- The torch.distributed retrieval setup created a synchronization bottleneck because rank 0 handled all worker index queries.
- Ray stateful actors allow multiple dedicated processes separate from training workers to load the index and process queries concurrently.
- In 4-GPU benchmarks with batch size 8 over 500 steps, four Ray retrieval processes lowered retrieval time from 3.438 to 1.66 seconds per call.

**Tags:** [LLMs](https://yomu.fyi/topic/llm), [Machine Learning](https://yomu.fyi/topic/machine-learning), [Performance](https://yomu.fyi/topic/performance), [Python](https://yomu.fyi/topic/python), [Scalability](https://yomu.fyi/topic/scalability)

[Read original post](https://huggingface.co/blog/ray-rag)