add_to_index uses too much GPU RAM and crashes

[MartinMashalov]

Hi there,

I am implementing Ragatouille at scale for a RAG model project. I am running the add_to_index() method sequentially but I keep running into memory issues with my GPU with 16 GB RAM. Does the add_to_index() method use GPU RAM or am I implementing something incorrectly?

Is there a way to optimize my GPU usage? Can I also run only on CPU?

Thanks.

Martin

Provided code:

from ragatouille import RAGPretrainedModel
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

RAG.index(
collection=[file_content],
document_metadatas=[{'source': 'initial_data.txt', 'page': 0}],
document_ids=[str(uuid.uuid4())],
index_name="test3",
split_documents=True,
max_document_length=300,
#use_faiss=True
)

for _ in range(10):
RAG.add_to_index(
new_collection=all_splits_rgt,
new_document_metadatas=rgt_metadatas,
new_document_ids=[str(uuid.uuid4()) for _ in range(len(all_splits_rgt))],
#index_name="test1",
#new_split_documents=True,
#new_max_document_length=300,
#use_faiss=True
)
retriever_advanced_text = RAG_redone.as_langchain_retriever(k=10)

[dineshdharme]

I am facing this issue even with just 100 documents trying to use just RAG.index().
Here's a simple code which reproduces the issue on 16GB VRAM.

from ragatouille import RAGPretrainedModel
from datasets import load_dataset

another_dataset = "kdawoud91/News_Articles"
df = load_dataset(another_dataset, cache_dir=cache_hf_dir)['train'].to_pandas()

df_limit = df[:100]
df_limit = df_limit.reset_index()
df_limit = df_limit.rename(columns={"index": "id"})
df_limit["id"] = df_limit["id"].map(repr)

list_news_articles = df_limit.text.values

RAG = RAGPretrainedModel.from_pretrained("jinaai/jina-colbert-v1-en")


RAG.index(
    collection=list_news_articles,
    index_name=INDEX_NAME,
    split_documents=True,
    use_faiss=True,
    max_document_length=4060,
)
 

[bclavie]

Hey, apologies for the very late answer. This seems to be an issue with the torch k-means implementation being far too hungry in some cases. Working on a fix now.

[santiagotri]

+1, I described my problem in #247

[matheusft]

I'm having this error when trying to add documents to an index created with RAGatouille (colbert-ir/colbertv2.0)

[Oct 11, 14:25:17] #> Optimizing IVF to store map from centroids to list of pids..
[Oct 11, 14:25:17] #> Building the emb2pid mapping..
Traceback (most recent call last):
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/ragatouille/RAGPretrainedModel.py", line 258, in add_to_index
    self.model.add_to_index(
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/ragatouille/models/colbert.py", line 179, in add_to_index
    self.model_index.add(
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/ragatouille/models/index.py", line 398, in add
    self.build(
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/ragatouille/models/index.py", line 243, in build
    indexer.index(name=index_name, collection=collection, overwrite=overwrite)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexer.py", line 80, in index
    self.__launch(collection)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexer.py", line 89, in __launch
    launcher.launch_without_fork(self.config, collection, shared_lists, shared_queues, self.verbose)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/infra/launcher.py", line 93, in launch_without_fork
    return_val = run_process_without_mp(self.callee, new_config, *args)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/infra/launcher.py", line 109, in run_process_without_mp
    return_val = callee(config, *args)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexing/collection_indexer.py", line 33, in encode
    encoder.run(shared_lists)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexing/collection_indexer.py", line 76, in run
    self.finalize() # Builds metadata and centroid to passage mapping
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexing/collection_indexer.py", line 396, in finalize
    self._build_ivf()
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexing/collection_indexer.py", line 482, in _build_ivf
    _, _ = optimize_ivf(ivf, ivf_lengths, self.config.index_path_)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexing/utils.py", line 13, in optimize_ivf
    all_doclens = load_doclens(index_path, flatten=False)
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexing/loaders.py", line 32, in load_doclens
    all_doclens = [ujson.load(open(filename)) for filename in doclens_filenames]
  File "/home/username/repo_name/.venv/lib/python3.10/site-packages/colbert/indexing/loaders.py", line 32, in <listcomp>
    all_doclens = [ujson.load(open(filename)) for filename in doclens_filenames]
ujson.JSONDecodeError: Expected object or value