langchain / / 2024. 11. 7. 08:45

reranker 사용 시 GPU가 필요할까?

RAG 품질개선 시 필수적인 요소인 reranker를 적용할 때 GPU가 반드시 필요할까?

RAG의 품질개선 기법 중 중요한 요소인 reranker를 적용하고 테스트를 해보았다. reranker는 BAAI/bge-reranker-v2-m3를 사용했다.

reranker가 cross-encoder 방식이어서 속도가 엄청 느리다고 하는데, 로컬개발환경에서는 생각보다는 빠르게 실행되어서 사용할만 했다. 하지만 개발서버에서 테스트를 해보니 너무 느렸다.

원인을 확인해보니 GPU 존재여부였다. 로컬환경은 Mac studio m1 max 였고 32코어 GPU가 있는 환경이고 서버환경은 GPU가 없었다. 속도차이가 너무 났다.

그래서 관련 자료를 찾아보니 누군가 블로그에 테스트한 결과를 잘 정리한 자료가 있어서 요약해 보았다.

[Speed Showdown for your RAG improvement: Reranker performance on CPU/GPU/TPU]

https://medium.com/@xiweizhou/speed-showdown-reranker-1f7987400077

from sentence_transformers import CrossEncoder
import time
# Load the models
base_model = CrossEncoder("BAAI/bge-reranker-base", trust_remote_code=True)
large_model = CrossEncoder("BAAI/bge-reranker-large", trust_remote_code=True)
v2m3_model = CrossEncoder("BAAI/bge-reranker-v2-m3", trust_remote_code=True)
# Function to rerank documents
def rerank_documents(model, query, contexts):
 start = time.time()
 rets = model.rank(query, contexts, batch_size=64)
 print(f"Ranking finished in {time.time()-start} seconds")
 return rets
# Sample query and contexts
query, contexts = ('Were Scott Derrickson and Ed Wood of the same nationality?',
 [{'id': 0, 'text': 'Scott Derrickson Scott Derrickson (born July 16, 1966) is an American director, screenwriter and producer. He lives in Los Angeles, California. He is best known for directing horror films such as "Sinister", "The Exorcism of Emily Rose", and "Deliver Us From Evil", as well as the 2016 Marvel Cinematic Universe installment, "Doctor Strange."'},
 {'id': 1, 'text': 'Ed Wood Edward Davis Wood Jr. (October 10, 1924 - December 10, 1978) was an American filmmaker, actor, writer, producer, and director.'}
 # 100 context in total ...
])
context_texts = [t['text'] for t in contexts]

# Run the tests
rerank_documents(base_model, query, context_texts)
rerank_documents(large_model, query, context_texts)
rerank_documents(v2m3_model, query, context_texts)

테스트 방법: 하나의 쿼리를 100개의 문서와 비교를 하는 시간을 측정한다. CPU 시간과 GPU 시간을 각각 측정한다.

bge-reranker-v2-m3의 결과는 아래와 같다.

[CPU]

Ranking finished in 257.12487530708313 seconds

[TPU]

Ranking finished in 13.545848369598389 seconds

[GPU T-4]

Ranking finished in 3.385152816772461 seconds

[GPU A-10]

Ranking finished in 1.404026746749878 seconds

위의 시간차이를 보면 알겠지만, GPU에서 테스트한 시간이 압도적으로 빠르다.

결론은 cross-encoder를 사용하는 reranker의 성능을 개선하려면 GPU는 반드시 필요할 것 같다.

출처

반응형
  • 네이버 블로그 공유
  • 네이버 밴드 공유
  • 페이스북 공유
  • 카카오스토리 공유