search

stanford-futuredata / ColBERT

ColBERT: state-of-the-art neural search (SIGIR'20, TACL'21, NeurIPS'21, NAACL'22, CIKM'22, ACL'23, EMNLP'23)
MIT License
2.95k stars 377 forks source link

Using GPU, ColBERT.try_load_torch_extensions from IndexUpdater reports "Error building extension 'segmented_lookup_cpp'" #276

Open sethhorrigan opened 10 months ago

sethhorrigan commented 10 months ago

We are working with ColBERT at https://github.com/stanford-futuredata/ColBERT/commit/89ab526a08db475e7843487bb87259630b5e4d7d, and we are seeing errors when attempting to use IndexUpdater on GPU.

1) Using only CPU Indexer.index succeeds and adding new texts via IndexUpdater.add also succeeds (expected) 2) Using GPU Indexer.index succeeds but adding new texts via IndexUpdater.add fails in the ninja call reporting Error building extension 'segmented_lookup_cpp' coming from the load call here

We see this using the libraries defined in https://github.com/c3-e/ColBERT/blob/89ab526a08db475e7843487bb87259630b5e4d7d/conda_env.yml

Is there perhaps a specific version of ninja we need in order for the IndexUpdater to initialize correctly? Or are there different versions of the CUDA libraries needed for the IndexUpdater to work correctly?

sethhorrigan commented 10 months ago

Looking into the implementation, it looks like IndexUpdater may be coded so it only ever uses CPU - https://github.com/c3-e/ColBERT/blob/f2235d1bdfc61c9894ddf1d16ac85ed089ed93bf/colbert/index_updater.py#L142