Commit f00f3e6
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: D1146857551 parent da3191e commit f00f3e6
7 files changed
Lines changed: 733 additions & 1284 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1123 | 1123 | | |
1124 | 1124 | | |
1125 | 1125 | | |
1126 | | - | |
1127 | | - | |
| 1126 | + | |
| 1127 | + | |
| 1128 | + | |
| 1129 | + | |
| 1130 | + | |
1128 | 1131 | | |
1129 | 1132 | | |
1130 | 1133 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
253 | 253 | | |
254 | 254 | | |
255 | 255 | | |
256 | | - | |
| 256 | + | |
| 257 | + | |
| 258 | + | |
| 259 | + | |
| 260 | + | |
| 261 | + | |
| 262 | + | |
| 263 | + | |
| 264 | + | |
| 265 | + | |
| 266 | + | |
257 | 267 | | |
258 | 268 | | |
259 | 269 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
78 | 78 | | |
79 | 79 | | |
80 | 80 | | |
| 81 | + | |
| 82 | + | |
| 83 | + | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
| 87 | + | |
| 88 | + | |
81 | 89 | | |
82 | 90 | | |
83 | 91 | | |
| |||
0 commit comments