AnswerDotAI / AnswerDotAI/RAGatouille
Cannot clear GPU memory when creating models with new indices
- 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