Skip to content

Commit cd22982

Browse files
Michael Norrisfacebook-github-bot
authored andcommitted
Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (facebookresearch#5500)
Summary: `GpuIndexCagra` had two multi-GPU build entry points totalling 16 positional arguments, neither of which fits the `Index` API. This removes one, folds the other into `train()`, and speeds up the build. **Usage.** Listing more than one device in the config selects the multi-GPU build; `train()` routes to it: ```python devices = faiss.Int32Vector() for i in range(8): devices.push_back(i) an = faiss.AllNeighborsCagraConfig() an.n_clusters = 16 # 0 = auto: max(2 * n_devices, 4) an.overlap_factor = 2 # must be >= 2 an.ivf_pq_search_batch_size = 8192 # 0 = cuVS default; caps IVF-PQ memory config = faiss.GpuIndexCagraConfig() config.graph_degree = 32 config.intermediate_graph_degree = 32 config.build_algo = faiss.graph_build_algo_IVF_PQ config.devices = devices # >1 device selects the multi-GPU build config.all_neighbors_params = an index = faiss.GpuIndexCagra(res, d, faiss.METRIC_L2, config) index.train(xb) # xb must stay alive until copyTo() completes cpu_index = faiss.IndexHNSWCagra() cpu_index.base_level_only = True index.copyTo(cpu_index) # required: the GPU index is not searchable on this path ``` Float32-only, does not copy `x`, and leaves `index_` empty so `copyTo()` is the only valid follow-up. Single-GPU behaviour is unchanged when `devices` is empty. **Deleted `trainMultiGpu`.** Sharded SNMG CAGRA produces per-shard graphs with zero cross-shard edges by construction, so it needs post-hoc stitching just to be usable, and it lost to the `all_neighbors` path on both build time and recall. It had no callers outside the benchmark and one test. **Folded `trainAllNeighbors` into `train()`.** Its 6 trailing bare scalars now live in the constructor-time config struct, which is the established Faiss GPU convention: | Old argument | Now | | --- | --- | | `devices` | `GpuIndexCagraConfig::devices`, also the dispatch predicate | | `build_algo` (0/1/2) | `GpuIndexCagraConfig::build_algo` | | `refinement_rate` | `AllNeighborsCagraConfig::refinement_rate` | | `n_clusters`, `overlap_factor`, `ivfpq_search_batch` | new `AllNeighborsCagraConfig` | This also kills a live footgun: the old `int build_algo` used an encoding (0=NN-descent, 1=brute-force, 2=IVF-PQ) that disagreed with the `graph_build_algo` enum in the same header (0=IVF_PQ, 1=NN_DESCENT). Both callers set `config.build_algo` and then passed an unrelated int, and the config field was silently dead. `BRUTE_FORCE` is appended to `graph_build_algo` (at the end, so existing values do not renumber) and the config field is now the single source of truth. **Graph pruning no longer copies the graph off the device.** Previously only detour counting ran on the GPU and pruning plus reverse-graph construction ran on the host, which was 35% of total build time. This part is **temporary and self-removing**. cuVS has no public API to prune a graph that is already on the device: its exported `helpers::optimize` takes host matrices, and the dispatch behind it erases the mdspan accessor to host memory, so cuVS's own device code path is unreachable from outside. Until that is fixed upstream, this reaches into cuVS's internal headers when the build has them available, and otherwise falls back to the public host API -- correct either way, just slower on the fallback. **Open-source builds get the fallback**, since cuVS does not install those headers; the `train()` docs say so. An upstream patch adding a `device_matrix_view` overload to `cuvs::neighbors::cagra::helpers::optimize` is prepared and will be submitted to rapidsai/cuvs. When it ships, the internal include, the build flag guarding it, and the fallback branch all get deleted and every build gets the fast path. At 50M vectors on 8 GPUs the optimize step goes from 119.1s to 1.65s (72x) and end-to-end build->serialize from 8.0 to 5.4 minutes, recall unchanged. **Collapsed the benchmark to one path.** With the stitching approaches gone, `bench_approaches.py` is a single-path tool for validating and tuning the production build, and reports a per-phase build breakdown plus an efSearch sweep of recall and QPS. Deliberately *not* merged into the config: - `ivf_pq_params` / `ivf_pq_search_params` are still not consulted on this path. cuVS derives them from the dataset shape and those derived values beat the static defaults in the faiss structs, so only the knobs cuVS cannot infer are overridden. - `faiss::cagra_build_algo` only has `{IVF_PQ, NN_DESCENT}`, so a `BRUTE_FORCE` config would silently degrade to NN-descent on the single-GPU path; `train_ex()` now rejects it there. `GpuIndexBinaryCagra` shares this config struct and has no multi-GPU build, so it rejects `devices.size() > 1` rather than silently building on one device. Behaviour changes: the default `build_algo` on the multi-GPU path is now `IVF_PQ` (the config default) rather than NN-descent (the old argument default), and `IndexHNSW::num_base_level_search_entrypoints` goes 32 -> 256, which measured better on both recall and QPS at every efSearch. Differential Revision: D114685755
1 parent 80a1656 commit cd22982

11 files changed

Lines changed: 657 additions & 1293 deletions

CMakeLists.txt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,10 @@ if(FAISS_ENABLE_CUVS AND NOT TARGET cuvs::cuvs)
125125
find_package(cuvs)
126126
endif()
127127

128+
if(FAISS_ENABLE_CUVS AND NOT TARGET rmm::rmm)
129+
find_package(rmm REQUIRED)
130+
endif()
131+
128132
add_subdirectory(faiss)
129133

130134
if(FAISS_ENABLE_GPU)

faiss/CMakeLists.txt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -448,6 +448,10 @@ endif()
448448

449449
add_library(faiss ${FAISS_SRC})
450450

451+
if(BUILD_SHARED_LIBS AND CMAKE_SYSTEM_NAME STREQUAL "Linux")
452+
target_link_options(faiss PRIVATE "LINKER:--no-undefined")
453+
endif()
454+
451455
add_library(faiss_avx2 ${FAISS_SRC})
452456
if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512_spr")
453457
set_target_properties(faiss_avx2 PROPERTIES EXCLUDE_FROM_ALL TRUE)

faiss/IndexHNSW.cpp

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
#include <limits>
1818
#include <memory>
1919
#include <queue>
20-
#include <random>
2120

2221
#include <cstdint>
2322
#include "faiss/Index.h"
@@ -1433,13 +1432,11 @@ void IndexHNSWCagra::search(
14331432
// first real candidate will always be strictly better.
14341433
nearest_d[i] = C::neutral();
14351434

1436-
std::random_device rd;
1437-
std::mt19937 gen(rd());
1438-
std::uniform_int_distribution<idx_t> distrib(
1439-
0, this->ntotal - 1);
1435+
// Seeded per query so entrypoints are reproducible.
1436+
SplitMix64RandomGenerator gen(i);
14401437

14411438
for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
1442-
auto idx = distrib(gen);
1439+
idx_t idx = gen.rand_int64() % this->ntotal;
14431440
auto distance = (*dis)(idx);
14441441
if (C::cmp(nearest_d[i], distance)) {
14451442
nearest[i] = static_cast<storage_idx_t>(idx);
@@ -1498,12 +1495,11 @@ void IndexHNSWCagra::range_search(
14981495
// real candidate will always be strictly better.
14991496
float nearest_d = C::neutral();
15001497

1501-
std::random_device rd;
1502-
std::mt19937 gen(rd());
1503-
std::uniform_int_distribution<idx_t> distrib(0, ntotal - 1);
1498+
// For reproducible entrypoint.
1499+
SplitMix64RandomGenerator gen(i);
15041500

15051501
for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
1506-
auto idx = distrib(gen);
1502+
idx_t idx = gen.rand_int64() % ntotal;
15071503
auto distance = (*dis)(idx);
15081504
// C::cmp(nearest_d, distance) is true iff distance is
15091505
// strictly better than the current nearest_d.

faiss/IndexHNSW.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,7 @@ struct IndexHNSWCagra : IndexHNSW {
253253
/// searches only the base level knn graph of the HNSW index.
254254
/// This parameter selects the entry point by randomly selecting
255255
/// some points and using the best one.
256-
int num_base_level_search_entrypoints = 32;
256+
int num_base_level_search_entrypoints = 256;
257257

258258
void add(idx_t n, const float* x) override;
259259

faiss/gpu/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -363,7 +363,7 @@ else()
363363

364364

365365
find_package(CUDAToolkit REQUIRED)
366-
target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs> $<$<BOOL:${FAISS_ENABLE_CUVS}>:OpenMP::OpenMP_CXX>)
366+
target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs> $<$<BOOL:${FAISS_ENABLE_CUVS}>:rmm::rmm> $<$<BOOL:${FAISS_ENABLE_CUVS}>:OpenMP::OpenMP_CXX>)
367367
target_compile_options(faiss_gpu_objs PRIVATE
368368
$<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all
369369
--expt-extended-lambda --expt-relaxed-constexpr

faiss/gpu/GpuIndexBinaryCagra.cu

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,13 @@ std::shared_ptr<GpuResources> GpuIndexBinaryCagra::getResources() {
7878
}
7979

8080
void GpuIndexBinaryCagra::train(idx_t n, const uint8_t* x) {
81+
// The config is shared with the float index; there is no binary
82+
// multi-GPU build, so reject rather than silently ignoring the request.
83+
FAISS_THROW_IF_MSG(
84+
cagraConfig_.devices.size() > 1,
85+
"binary CAGRA has no multi-GPU build; "
86+
"GpuIndexCagraConfig::devices must name at most one device");
87+
8188
DeviceScope scope(cagraConfig_.device);
8289
if (this->is_trained) {
8390
FAISS_ASSERT(index_);

0 commit comments

Comments
 (0)