Retrieval Augmented Generation with Huggingface Transformers and Ray

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.

Takeaways

1. 1

   The torch.distributed retrieval setup created a synchronization bottleneck because rank 0 handled all worker index queries.

2. 2

   Ray stateful actors allow multiple dedicated processes separate from training workers to load the index and process queries concurrently.

3. 3

   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.