AnswerDotAI / AnswerDotAI/RAGatouille

Cannot clear GPU memory when creating models with new indices

Open
#270 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
4k
Forks
276
PR merge metrics
No merged PRs in 30d

Description

Hi! I have a query that I want to run over individual documents, each with their own index. There are too many documents to put into memory at the same time, and so I would like to clear CUDA memory between each document search. The RAG model for each document is either initiated with RAGPretrainedModel.from_pretrained (if the doc hasn't been indexed before), or with RAGPretrainedModel.from_index (if the doc has been previously indexed).

Below is pseudocode of the document search loop:

```
model_path = ... # same model for all documents
doc_ids = [...] # list of docs
index_paths = {doc_id: make_index_path(doc_id) for doc_id in doc_ids} # output locations of RAG.index for each doc_id
results = dict()
for doc_id, index_path in index_paths.items():
document = get_document(doc_id)
if not os.path.exists(index_path):
RAG = RAGPretrainedModel.from_pretrained(model_path, index_root=index_root)

print(f"Indexing document {doc_id}...")
print(f"len(document {doc_id}): {len(document)}")
RAG.index(
collection=[d['content'] for d in document],
document_ids=[d['document_id'] for d in document],
index_name=f"doc_{doc_id}",
max_document_length=512,
split_documents=False,
use_faiss=True,
)
else:
print(f"Index for document {doc_id} already exists")
RAG = RAGPretrainedModel.from_index(index_path)

results[doc_id] = RAG.search(
query=query,
index_name=f"doc_{doc_id}",
k=search_k,
)

# WANT TO DELETE THIS RAG MODEL AND CLEAR GPU MEMORY HERE
...
```

I can't figure out a way to clear the RAGPretrainedModels' memory between documents. Deleting the RAG model and calling torch.cuda.empty_cache() doesn't release anything. Is there a way to delete the model and release everything it's keeping in GPU without having to kill the process?

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the RAGPretrainedModel.from_pretrained and from_index entry points, then trace the index and search calls in the provided loop. Compare deleting the model and calling torch.cuda.empty_cache() between documents; done means the model releases its GPU memory without requiring the process to be killed.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.