Skip to content

Commit f00f3e6

Browse files
Michael Norrisfacebook-github-bot
authored andcommitted
Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#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 moves the graph-pruning phase fully onto the GPU. **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 # do not lower; see below an.ivf_pq_search_batch_size = 8192 # 0 = cuVS default; caps IVF-PQ workspace config = faiss.GpuIndexCagraConfig() config.graph_degree = 32 config.intermediate_graph_degree = 32 config.build_algo = faiss.graph_build_algo_IVF_PQ config.refine_rate = 2.0 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 ``` That path is 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` | `GpuIndexCagraConfig::refine_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 now runs entirely on the GPU.** The previous implementation offloaded only detour counting and did pruning and reverse-graph construction on the host, which was 35% of total build time. It is replaced by a direct call to the cuVS implementation with device-resident mdspans: ```cpp cuvs::neighbors::cagra::detail::graph::optimize<uint32_t>( single_gpu_res, d_knn.view(), d_cagra.view(), cagraConfig_.guarantee_connectivity); ``` This deliberately bypasses `cuvs::neighbors::cagra::optimize()`, whose dispatch erases the mdspan accessor to `raft::memory_type::host`. That makes the device-resident branch of `make_reverse_graph_gpu` unreachable and degrades the reverse-graph phase into `graph_degree` separate host gathers, each with its own H2D copy and a full stream synchronisation. Passing device mdspans keeps prune, reverse graph and merge on device. The int64 -> uint32 narrowing that `all_neighbors` output requires is now a device kernel rather than a D2H copy plus host loop. At 50M vectors this takes graph optimize from 119.1s to 1.65s (72x) and end-to-end build->serialize from 8.0 to 5.4 minutes, with recall unchanged. `AllNeighborsCagraConfig::multi_gpu_optimize` is removed: it selected the host-side implementation that no longer exists. **Collapsed the benchmark to one path.** With the stitching approaches gone, `bench_approaches.py` is now a single-path tool for validating and tuning the production build: no `--approaches` flag, no per-approach labelling, no `IndexShards` handling in the eval helpers. It gains `--guarantee-connectivity` to control the MST pass in `cagra::optimize`. Deliberately *not* merged into the config: - `ivf_pq_params` / `ivf_pq_search_params` are still not consulted on this path. cuVS derives `n_lists`, `pq_dim` and the kmeans trainset fraction from the dataset shape (`n_lists = n/2000`, i.e. 50000 at 100M vectors, versus the static default of 1024). Applying `IVFPQ*CagraConfig` wholesale would discard that tuning. The IVF-PQ search batch cap therefore keeps its own field, with 0 meaning "leave cuVS's dataset-derived default". - `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 change: the default `build_algo` on the multi-GPU path is now `IVF_PQ` (the config default) rather than NN-descent (the old argument default). That matches the single-GPU path and the recommended large-scale config. Differential Revision: D114685755
1 parent da3191e commit f00f3e6

7 files changed

Lines changed: 733 additions & 1284 deletions

File tree

faiss/IndexHNSW.cpp

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1123,8 +1123,11 @@ void IndexHNSWCagra::search(
11231123
// first real candidate will always be strictly better.
11241124
nearest_d[i] = C::neutral();
11251125

1126-
std::random_device rd;
1127-
std::mt19937 gen(rd());
1126+
// Seeded per query rather than from random_device: seeding a
1127+
// mt19937 costs more than the handful of samples drawn from
1128+
// it, and a nondeterministic entry point makes recall
1129+
// irreproducible run to run.
1130+
std::mt19937 gen(0x9e3779b9u ^ (uint32_t)i);
11281131
std::uniform_int_distribution<idx_t> distrib(
11291132
0, this->ntotal - 1);
11301133

faiss/IndexHNSW.h

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,17 @@ 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+
///
257+
/// Each sample costs one full distance computation before the beam search
258+
/// starts, so this is a fixed per-query tax that does not shrink with
259+
/// efSearch. Up to a point a better entry point makes the beam search
260+
/// converge in fewer hops and more than pays for itself: measured on 100M
261+
/// x 129d, going from 32 to 256 improved both recall and QPS at every
262+
/// efSearch. Beyond that it stops paying -- the nearest of N uniform
263+
/// samples improves only as N^(-1/d), so 4096 samples (128x the work of
264+
/// 32) bought just +0.012 recall while costing 19% QPS at efSearch=512 and
265+
/// 81% at efSearch=16. Raise it only if you serve at high efSearch.
266+
int num_base_level_search_entrypoints = 256;
257267

258268
void add(idx_t n, const float* x) override;
259269

faiss/gpu/GpuIndexBinaryCagra.cu

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

8080
void GpuIndexBinaryCagra::train(idx_t n, const uint8_t* x) {
81+
// GpuIndexCagraConfig is shared with the float index, whose multi-GPU
82+
// build path has no binary equivalent. Reject rather than silently
83+
// building on one device.
84+
FAISS_THROW_IF_MSG(
85+
cagraConfig_.devices.size() > 1,
86+
"binary CAGRA has no multi-GPU build; "
87+
"GpuIndexCagraConfig::devices must name at most one device");
88+
8189
DeviceScope scope(cagraConfig_.device);
8290
if (this->is_trained) {
8391
FAISS_ASSERT(index_);

0 commit comments

Comments
 (0)