CUDA RAM OOM for large query lengths and document lengths = 4096

[thakur-nandan]

Hi @NohTow,

Thanks for this amazing package. I've been trying to evaluate GTE-ModernColBERT-v1 on FreshStack.
I used the fastPlaid script, but since my max lengths are quite large in FreshStack (4096 for both query and document), I encounter a OOM even with a single H200 GPU with ~140GB VRAM.

Since you have better experience, what's the possible solution here?
Should I reduce my max length to 512 and evaluate which fits within the H200 GPU, or should I use the CPU for evaluation instead?

The code is bug-free: https://github.com/fresh-stack/freshstack/blob/main/examples/evaluation/pylate_retrieval_evaluation.py

Thanks,
Nandan Thakur

[NohTow]

Hey Nandan!
Thanks for the kind words!

What a coïncidence, I am also working on evaluating models on code benchmarks (including FreshStack!) and I've ran into the same issue (that's also why I did not report FreshStack earlier).

I am still using OG PLAID (because fast-plaid has still some high memory consumption when creating indexes, but here it's not the issue) and what makes it crash is the creating of the big stridded tensor that is like length_longest_doc * num_reranking_doc * embedding_dim, which is just enormous when considering long document. Interestingly, it did not happen when I was evaluating LongEmbed, prolly because there were less candidate per query.

We are actively working on solving this, @raphaelsty for example shipped fast-plaid 1.3 which reduced memory usage, but I still face the issue to some extend (so you can try upgrading, make sure to select the version with the correct torch version, but I do not think it'll solve your issue totally). He has plan to make fast-plaid streaming friendly which should solve the issue altogether.

In the mean time, the only things I have managed to tweak to be able to run some evals on FreshStack is these two parameters:

batch_size
n_full_scores

The first one can be lowered without any change to the result, it'll just make the search slower, but it can still crash if the stridded document is too big even for one query at a time. Lowering n_full_scores allows to lower the memory consumption, but it can also change the results because we are considering less documents... I unfortunately have no other solutions until we fix the root cause!

Edit:
You can also indeed try to run the search on CPU to avoid cuda OOM
You can set device="cpu" when initing the fast-plaid index, this way you can still use the GPUs for encoding the docs/queries, I am not too sure about the speed though, our nodes were melting today

[NohTow]

Hello,
As described there, we did a lot of improvement to memory usage and I managed to evaluate on FreshStack with reasonably large lengths (but the results of ColBERT models did not seems insane). You might have to lower the batch sizes but it should be runnable. Let me know if you are able to run it and if not, whether it crashes during indexation or search.
Would be happy to try to understand why ColBERT models doesn't seem insane on FreshStack, including the code ones. Happy to debug!