Skip to content

fix algorithm of spreading vectors over shards #3374

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
simshi wants to merge 2 commits into from
Closed

fix algorithm of spreading vectors over shards #3374

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7adc56b
fix algorithm of spreading vectors over shards
simshi Apr 19, 2024
8143e0a
Merge branch 'main' into shardspread
junjieqi May 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions faiss/gpu/GpuIcmEncoder.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ void GpuIcmEncoder::encode(
size_t n,
size_t ils_iters) const {
size_t nshards = shards->size();
size_t shard_size = (n + nshards - 1) / nshards;
size_t base_shard_size = n / nshards;
Copy link
Contributor

@junjieqi junjieqi Apr 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is it (n / nshards) or (n - 1)/ nshards?

Copy link
Contributor Author

@simshi simshi May 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

integer computing, round down already. remain = n - nshards * base_shard_size, then 0 <= remain < nshards, so each shard takes at least base_shard_size vectors, and the first remain shards takes one extra vector each.


auto codebooks = lsq->codebooks.data();
auto M = lsq->M;
Expand All @@ -94,8 +94,14 @@ void GpuIcmEncoder::encode(

// split input data
auto fn = [=](int idx, IcmEncoderImpl* encoder) {
size_t i0 = idx * shard_size;
size_t ni = std::min(shard_size, n - i0);
size_t i0 = idx * base_shard_size + std::min(size_t(idx), n % nshards);
Copy link
Contributor

@junjieqi junjieqi Apr 29, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is it std::min(size_t(idx), n % nshards)? Btw, should be (n-1) % nshards?

Copy link
Contributor Author

@simshi simshi May 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

as my previous comment, you can verify it over the case n % nshards == 0, then (n -1) % nshards is meaningless

size_t ni = base_shard_size;
if (ni < n % nshards) {
++ni;
}
if (ni <= 0) { // only if n < nshards
return;
}
auto xi = x + i0 * d;
auto ci = codes + i0 * M;
std::mt19937 geni(idx + seed); // different seed for each shard
Expand Down