Huggingface가 RAG 모델의 문서 검색 메커니즘을 torch.distributed에서 Ray로 교체해 검색 속도를 2배 향상 및 다중 GPU 학습 확장성 개선

Context

RAG(Retrieval Augmented Generation) 모델의 분산 학습 시 torch.distributed 기반 문서 검색이 병목 현상을 야기했다. Rank 0 워커가 모든 워커의 입력을 수신하고 인덱스 쿼리를 수행한 후 결과를 다시 전송하는 동기화 구조로 인해 GPU 수 증가에 따라 검색 레이턴시가 2.12초에서 3.438초로 증가했다. 추가로 torch.distributed 구현은 PyTorch에 종속적이고 문서 인덱스 크기가 커 각 워커마다 복제 로드가 불가능한 한계가 있었다.

Technical Solution

• 문서 검색 프로세스 분리: torch.distributed의 단일 Rank 0 워커 구조에서 Ray의 상태 저장 액터(stateful actor) 기반 다중 프로세스 구조로 변경
• 인덱스 로드 분산: 대규모 문서 인덱스를 학습 프로세스와 독립적인 Ray 프로세스들이 로드하고 관리
• 동시 검색 처리: 단일 검색 프로세스의 병목 제거를 위해 여러 Ray 검색 프로세스를 병렬로 운영
• 프레임워크 의존성 제거: PyTorch 프로세스 그룹 종속성을 제거해 프레임워크 독립적 검색 구현 실현
• 통합 인터페이스 제공: Huggingface의 PyTorch Lightning 기반 파인튜닝 스크립트에 --distributed_retriever ray 옵션 추가

Impact

• 2 GPU 학습 시: torch.distributed 2.12초/검색 → Ray(2 프로세스) 1.49초/검색 (29.7% 개선)
• 4 GPU 학습 시: torch.distributed 3.438초/검색 → Ray(4 프로세스) 1.66초/검색 (51.7% 개선)
• Ray 검색 프로세스 수 증가 시 성능 향상: 2 프로세스 대비 4 프로세스 사용 시 2 GPU에서 1.49초 → 1.145초, 4 GPU에서 2.029초 → 1.66초

Key Takeaway

분산 시스템에서 특정 프레임워크의 통신 레이어에 종속된 기능은 전문화된 분산 프로그래밍 라이브러리로 분리하면 성능과 유연성을 동시에 확보할 수 있다. 액터 기반 병렬 처리는 학습과 데이터 검색 같은 서로 다른 리소스 특성의 작업을 독립적으로 스케일링하는 데 효과적이다.