Error while indexing large datasets like webis-touche using Fast Plaid

[meetdoshi90]

Error message on torch quantile with FastPLAID
I was trying to evaluate web-touche using the evaluation code provided by pylate (1.3.4). I am using FastPLAID index with torch 2.8.0. Any wordarounds on this would be appreciated. Works fine for other smaller datasets.
Related issues: torch-64947

thread '' panicked at /root/.cache/puccinialin/cargo/registry/src/index.crates.io-1949cf8c6b5b557f/tch-0.20.0/src/wrappers/tensor_generated.rs:14691:57:
called Result::unwrap() on an Err value: Torch("quantile() input tensor is too large\nException raised from quantile_compute at /pytorch/aten/src/ATen/native/Sorting.cpp:290 (most recent call first):\nframe #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits, std::allocator >) + 0x80 (0x1486a5c18eb0 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libc10.so)\nframe #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, char const*) + 0x69 (0x1486a5bb5b5f in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libc10.so)\nframe #2: + 0x1c8da77 (0x1486e35bea77 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #3: + 0x1c8ed4b (0x1486e35bfd4b in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #4: at::native::quantile(at::Tensor const&, at::Tensor const&, std::optional, bool, std::basic_string_view<char, std::char_traits >) + 0x57 (0x1486e35c02b7 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #5: + 0x2be537d (0x1486e451637d in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #6: at::_ops::quantile::call(at::Tensor const&, at::Tensor const&, std::optional, bool, std::basic_string_view<char, std::char_traits >) + 0x228 (0x1486e3d3cc78 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #7: + 0x1e6924 (0x148376870924 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #8: + 0x161e1f (0x1483767ebe1f in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #9: + 0x191486 (0x14837681b486 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #10: + 0x191f52 (0x14837681bf52 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #11: + 0x190050 (0x14837681a050 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #12: + 0x191860 (0x14837681b860 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\n\nframe #28: + 0x295d0 (0x1487215b05d0 in /lib64/libc.so.6)\nframe #29: __libc_start_main + 0x80 (0x1487215b0680 in /lib64/libc.so.6)\n")
note: run with RUST_BACKTRACE=1 environment variable to display a backtrace
Tasks: 0%| | 0/5 [1:38:14<?, ?it/s]
Traceback (most recent call last):
File "/home/embedding/meet/TableRetrieval/Table-RAG/finetuning/granite-embeddings/scripts/landmark_index_colbert.py", line 519, in
index.add_documents(
File "/home/meet/new_conda/anaconda3/lib/python3.12/site-packages/pylate/indexes/plaid.py", line 217, in add_documents
return self._index.add_documents(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/meet/new_conda/anaconda3/lib/python3.12/site-packages/pylate/indexes/fast_plaid.py", line 189, in add_documents
self.fast_plaid.create(
File "/home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/search/fast_plaid.py", line 283, in create
fast_plaid_rust.create(
pyo3_runtime.PanicException: called Result::unwrap() on an Err value: Torch("quantile() input tensor is too large\nException raised from quantile_compute at /pytorch/aten/src/ATen/native/Sorting.cpp:290 (most recent call first):\nframe #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits, std::allocator >) + 0x80 (0x1486a5c18eb0 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libc10.so)\nframe #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, char const*) + 0x69 (0x1486a5bb5b5f in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libc10.so)\nframe #2: + 0x1c8da77 (0x1486e35bea77 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #3: + 0x1c8ed4b (0x1486e35bfd4b in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #4: at::native::quantile(at::Tensor const&, at::Tensor const&, std::optional, bool, std::basic_string_view<char, std::char_traits >) + 0x57 (0x1486e35c02b7 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #5: + 0x2be537d (0x1486e451637d in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #6: at::_ops::quantile::call(at::Tensor const&, at::Tensor const&, std::optional, bool, std::basic_string_view<char, std::char_traits >) + 0x228 (0x1486e3d3cc78 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/torch/lib/libtorch_cpu.so)\nframe #7: + 0x1e6924 (0x148376870924 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #8: + 0x161e1f (0x1483767ebe1f in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #9: + 0x191486 (0x14837681b486 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #10: + 0x191f52 (0x14837681bf52 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #11: + 0x190050 (0x14837681a050 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\nframe #12: + 0x191860 (0x14837681b860 in /home/meet/new_conda/anaconda3/lib/python3.12/site-packages/fast_plaid/fast_plaid_rust.cpython-312-x86_64-linux-gnu.so)\n\nframe #28: + 0x295d0 (0x1487215b05d0 in /lib64/libc.so.6)\nframe #29: __libc_start_main + 0x80 (0x1487215b0680 in /lib64/libc.so.6)\n")

[NohTow]

Thanks for reporting (also with the related PyTorch issue)!
@raphaelsty is currently working hard on making Fast-Plaid more scalable (also w.r.t memory usage), so he'll surely have a look at this as well! (maybe opening the same issue but on the fast-plaid repo would be great for tracking)

In the mean time, the workaround within PyLate would be to use the og PLAID by setting use_fast=False when creating the PLAID index. It's slower, but it handles very large BEIR dataset (including MS MARCO) without any issue, so until Raph fixes everything, this is the easiest/quickest workaround

[meetdoshi90]

Thanks @NohTow, I will open an issue. Yes, Stanford PLAID works correctly. I have a use case and wanted to benchmark on FastPLAID.

[meetdoshi90]

Edit: This also fails on Stanford PLAID due to torch quantile limitation:
Need to make this 50k a configurable parameter.

For more details, refer to this issue: fastplaid/27

[NohTow]

Edit: This also fails on Stanford PLAID due to torch quantile limitation:

Uh, this is a bit odd because I ran the whole BEIR suite multiple time, including webis (and larger dataset) without facing this issue.
Which Torch version are you using?

[meetdoshi90]

@NohTow, I think it worked because the embedding dimension was lower than 335 for the ColBERT that you used. If you try with a larger embedding, it will crash. I am using torch 2.8.0.

[NohTow]

Oh yeah that's correct my bad I dismissed the embedding dim because we are usually using 128 (or lower).
So indeed, making the maximum value configurable is a solution, for example by making sure we cap not by the number of heldout samples but by num_heldout_samples * embedding_dim (and put max to 128 * 50k), that being said I wonder if this is the optimal thing to do, because it means you'll be using a pretty small number of samples to compute the codebook.
Maybe we could replace the quantile from torch by something better instead? (My main worry is limiting edge cases/effects)

[meetdoshi90]

I can work on PR based on what change is preferred. If you want to replace torch quantile, it will require changes in rust implementation as well. There are a couple of reference implementations available in the torch issue. I can try and benchmark those if you like.

[NohTow]

I mean, if we can have the same behavior because this is just a limitation of torch, I am fine with this, maybe @raphaelsty will have to redo some work for FastPlaid in Rust, but I am a bit scared to lower the number of elem to compute quant.
Maybe I could run some benches with lower values, but if this is just a torch issue then maybe a fix is better anyways no?

[meetdoshi90]

Agreed. It doesn't look like it will be fixed in torch anytime soon. I will try it out and get back.

[NohTow]

Would be awesome, thank you!