search

terrierteam / pyterrier_colbert

82 stars 35 forks source link

How to run colbert indexing in multi-GPUs #54

Open cramraj8 opened 1 year ago

cramraj8 commented 1 year ago

I tried to run colbert indexing on trec-deep-learning-passages, and my environment is available with 4 GPUs. But when I call the APIs for indexing such as below it only utilizing 1 GPU.

from pyterrier_colbert.indexing import ColBERTIndexer
indexer = ColBERTIndexer(checkpoint, "/path/to/index", "index_name", ids=True)
indexer.index(dataset.get_corpus_iter())

How can I leverage all 4 GPUs with pyterrier API ?

cmacdonald commented 1 year ago

Good question @cramraj8 ! (cc @seanmacavaney)

Our setup is driven by Docker environments with primarily single GPU access. I know the underlying Colbert codebase can do distributed indexing, but we haven't integrated that.

Can you print the contents of from colbert.parameters import DEVICE and tell us which device is associated with indexer.colbert?

The relevant actual model encoding line is https://github.com/cmacdonald/ColBERT/blob/v0.2/colbert/modeling/inference.py#L30

I wonder if you can wrap indexer.colbert in torch.nn.DataParallel, as per https://www.run.ai/guides/multi-gpu/pytorch-multi-gpu-4-techniques-explained#Technique-1 and https://stackoverflow.com/a/64825728. Probably you can increase the value of indexer.args.bsize if you do?

Please let us know how you get on.

Craig

cmacdonald commented 1 year ago

any update @cramraj8