Hugging Face Transformers와 Ray를 활용한 Retrieval Augmented Generation

Hugging Face는 Retrieval Augmented Generation (RAG) 모델의 문서 검색 메커니즘에 Ray를 통합했습니다. 이 통합을 통해 검색 호출 지연 시간을 최대 2배까지 줄이고, 지식 집약적인 NLP 작업에 대한 분산 미세 조정(fine-tuning)의 확장성을 개선했습니다.

Retrieval Augmented Generation (RAG)의 이해

Retrieval Augmented Generation (RAG)은 실행 중에 Wikipedia 텍스트 코퍼스와 같은 외부 지식 베이스에서 문맥적 문서를 검색하여 지식을 강화하는 sequence-to-sequence (seq2seq) 아키텍처입니다. 모델에 내장된 파라미터에만 의존하는 표준 모델과 달리, RAG는 이러한 내부 파라미터와 외부 구절에서 검색된 정보를 결합하여 출력을 생성합니다. 이러한 이중 소스 접근 방식 덕분에 RAG는 질의응답과 같은 지식 집약적인 작업에서 다른 최첨단 모델보다 뛰어난 성능을 발휘합니다.

Ray를 통한 분산 미세 조정의 확장

검색 단계는 RAG의 성능에 필수적이지만, 분산 미세 조정 과정에서 상당한 복잡성을 초래합니다. 데이터 병렬 학습 루틴에서 문서 인덱스는 각 학습 워커(worker)가 복제된 사본을 로드하기에는 너무 큰 경우가 많아 잠재적인 병목 현상을 생성합니다.

이전에는 RAG 미세 조정을 위해 문서 검색에 torch.distributed 통신 패키지를 사용했습니다. 그러나 이 구현에는 두 가지 주요 제한 사항이 있었습니다:

  1. 동기화 병목 현상: rank 0 워커가 모든 워커로부터 입력을 받고, 인덱스 쿼리를 수행하고, 결과를 다시 배포하는 책임을 맡았기 때문에 학습 워커 수가 증가함에 따라 성능이 제한되었습니다.
  2. 프레임워크 의존성: 검색 프로세스 그룹이 학습 프로세스 그룹에 묶여 있어, 학습 프로세스에 PyTorch를 사용해야만 했습니다.

torch.distributed를 범용 분산 및 병렬 프로그래밍을 위한 Python 라이브러리인 Ray로 교체함으로써, Hugging Face는 프레임워크에 구애받지 않는 구현을 만들었습니다. Ray의 상태 유지(stateful) 액터 추상화를 사용하면 학습 프로세스와 분리된 여러 프로세스가 인덱스를 로드하고 검색 쿼리를 동시에 처리할 수 있어 rank 0 병목 현상을 제거할 수 있습니다.

성능 벤치마크

Ray를 통합하면 torch.distributed 구현과 비교하여 멀티 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

참고: 벤치마크는 GPU당 배치 크기 8로 500회의 학습 스텝 동안 수행되었으며, rank 0 워커에서 문맥적 문서를 검색하는 데 걸리는 시간을 측정했습니다.

구현 및 사용법

사용자는 Hugging Face가 제공하는 PyTorch Lightning 기반 미세 조정 스크립트를 사용하여 Ray 기반 검색을 구현할 수 있습니다. 프로세스는 ray와 transformers를 설치하고, 미세 조정 스크립트에서 다음 설정을 실행하는 것을 포함합니다:

  • Distributed Retriever: ray로 설정.
  • Retrieval Workers: --num_retrieval_workers 플래그를 통해 지정.

더 나은 최적화를 원하는 경우, 모델 정확도를 높이기 위해 확장 가능한 하이퍼파라미터 튜닝을 위해 Ray Tune과의 통합을 사용할 수 있습니다.

Sources