RTX 3090에서 RAG 모델 실행 시 CUDA 사용하기

최근 많은 AI 연구자들과 개발자들이 RAG(Retrieval-Augmented Generation) 모델을 활용하여 자연어 처리(NLP) 작업을 수행하고 있습니다. 그러나 기본적으로 PyTorch는 Mac에서는 MPS를, 일반적인 GPU 환경에서는 CUDA를 사용하도록 설정해야 합니다. 특히, RTX 3090과 같은 강력한 GPU를 활용하는 경우 CUDA 환경에서 최적의 성능을 내도록 설정하는 것이 중요합니다.

이 글에서는 RTX 3090을 활용하여 RAG 모델을 CUDA 환경에서 실행하는 방법과 발생할 수 있는 오류를 해결하는 방법을 설명하겠습니다.

1. 기본 코드 및 CUDA 설정

먼저, RAG 모델을 실행할 기본 코드를 살펴보겠습니다. 이 코드에서는 facebook/rag-sequence-nq 모델을 사용하여 질문에 대한 답변을 생성합니다. 아래 코드는 기존 MPS 설정이 되어 있던 부분을 CUDA로 변경한 버전입니다.

from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
from datasets import load_dataset
import torch

# CUDA 사용 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# wiki_dpr 데이터셋을 다운로드 (trust_remote_code=True 추가)
dataset = load_dataset("wiki_dpr", "psgs_w100.nq.exact", trust_remote_code=True)

# RAG 모델과 토크나이저 로드
tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq")
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq").to(device)

# 질문 입력
question = "What is the capital of France?"

# 질문을 토큰화
inputs = tokenizer(question, return_tensors="pt").to(device)

# 관련 문서 검색
retrieved_docs = retriever.retrieve(question)

# RAG 모델을 통해 답변 생성
generated_ids = model.generate(**inputs, retriever=retrieved_docs)

# 생성된 답변 디코딩
answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

print("Answer:", answer[0])

주요 변경 사항

• torch.device("mps" if torch.backends.mps.is_available() else "cpu") 대신 torch.device("cuda" if torch.cuda.is_available() else "cpu")를 사용하여 CUDA가 활성화된 경우 GPU를 활용하도록 수정했습니다.
• load_dataset("wiki_dpr", "psgs_w100.nq.exact", trust_remote_code=True)를 추가하여 데이터셋 로드 중 발생하는 오류를 방지했습니다.

2. 실행 중 발생할 수 있는 오류와 해결 방법

위 코드를 실행하는 과정에서 여러 가지 오류가 발생할 수 있습니다. 대표적인 오류와 해결 방법을 정리해 보았습니다.

(1) trust_remote_code=True 관련 오류

오류 메시지

ValueError: The repository for wiki_dpr contains custom code which must be executed to correctly load the dataset.
Please pass the argument `trust_remote_code=True` to allow custom code to be run.

해결 방법

load_dataset을 호출할 때 trust_remote_code=True를 추가하면 해결됩니다.

dataset = load_dataset("wiki_dpr", "psgs_w100.nq.exact", trust_remote_code=True)

(2) CUDA가 정상적으로 작동하지 않는 경우

오류 메시지

RuntimeError: CUDA out of memory. Tried to allocate ...

해결 방법

• 실행 중인 다른 GPU 프로세스를 종료합니다.
• nvidia-smi 명령어를 사용해 GPU 메모리 상태를 확인합니다.
• 모델을 half precision (float16)으로 변환하여 메모리 사용량을 줄일 수 있습니다.

model = model.half()

3. 성능 최적화

RTX 3090을 최대로 활용하려면 몇 가지 추가적인 최적화 방법을 고려할 수 있습니다.

(1) torch.cuda.amp를 활용한 Mixed Precision Training

torch.cuda.amp를 사용하면 자동으로 연산 속도를 최적화할 수 있습니다.

from torch.cuda.amp import autocast

with autocast():
    generated_ids = model.generate(**inputs, retriever=retrieved_docs)

(2) 배치 크기 조절

배치 크기를 줄이면 GPU 메모리 사용량을 줄일 수 있습니다.

inputs = tokenizer(question, return_tensors="pt", max_length=512, truncation=True).to(device)

결론

RTX 3090을 활용해 RAG 모델을 실행하려면 CUDA 환경을 설정하고, 데이터셋 로드 시 trust_remote_code=True 옵션을 추가해야 합니다. 또한, GPU 메모리 문제를 해결하고 성능을 최적화하는 다양한 방법을 적용하면 더욱 효율적으로 모델을 사용할 수 있습니다.

앞으로도 AI 모델을 실행하는 과정에서 발생할 수 있는 다양한 문제와 해결책을 공유할 예정이니 많은 관심 부탁드립니다!