Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500) by mnorris11 · Pull Request #5500 · facebookresearch/faiss

mnorris11 · 2026-08-04T03:58:28Z

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:

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

meta-codesync · 2026-08-04T03:58:35Z

@mnorris11 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114685755.

…okresearch#5500) Summary: `GpuIndexCagra` had two multi-GPU build entry points totalling 16 positional arguments, neither of which fits the `Index` API. This removes one and folds the other into `train()`. **Delete `trainMultiGpu` (approaches B and C).** Sharded SNMG CAGRA produces per-shard graphs with zero cross-shard edges by construction, so it needs post-hoc stitching (GPU brute-force or CPU HNSW) 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. The benchmark's approaches B and C go with it; approach A (`IndexShards`) stays, since it does not use `trainMultiGpu` and is still a useful baseline. **Fold `trainAllNeighbors` into `train()`.** It was a 9-argument method with 6 trailing bare scalars: ``` index.trainAllNeighbors(n, ptr, devices, 0, 0, True, 2, 2.0, 8192) ``` `train_ex()` now dispatches to the private `trainAllNeighbors_()` when `GpuIndexCagraConfig::devices` names more than one device, which is the established Faiss GPU convention (build knobs live in the constructor-time config struct). The six scalars resolve as: | 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`, `multi_gpu_optimize`, `ivfpq_search_batch` | new `AllNeighborsCagraConfig` | This also kills a live footgun: `trainAllNeighbors` took an `int build_algo` whose encoding (0=NN-descent, 1=brute-force, 2=IVF-PQ) 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 on that path. `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. Notes on what deliberately did NOT get 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, 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. Behavior 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 100M config. Differential Revision: D114685755

…okresearch#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

…okresearch#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

meta-codesync · 2026-08-14T23:10:00Z

This pull request has been merged in 6644dcf.

facebook-github-tools · 2026-08-15T01:16:04Z

This pull request has been reverted by 1f93154.

meta-cla Bot added the CLA Signed label Aug 4, 2026

meta-codesync Bot added the meta-exported label Aug 4, 2026

meta-codesync Bot changed the title ~~Fold multi-GPU CAGRA build into train(), delete trainMultiGpu~~ Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500) Aug 4, 2026

mnorris11 force-pushed the export-D114685755 branch from fa19d6e to 303dc3c Compare August 4, 2026 04:13

meta-codesync Bot changed the title ~~Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500)~~ Fold multi-GPU CAGRA build into train(), delete trainMultiGpu Aug 4, 2026

mnorris11 force-pushed the export-D114685755 branch 2 times, most recently from e7d6e15 to a54a9a3 Compare August 6, 2026 04:17

meta-codesync Bot changed the title ~~Fold multi-GPU CAGRA build into train(), delete trainMultiGpu~~ Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500) Aug 6, 2026

mnorris11 force-pushed the export-D114685755 branch from a54a9a3 to f00f3e6 Compare August 6, 2026 04:36

meta-codesync Bot changed the title ~~Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500)~~ Fold multi-GPU CAGRA build into train(), delete trainMultiGpu Aug 6, 2026

mnorris11 force-pushed the export-D114685755 branch from f00f3e6 to 558eb9c Compare August 6, 2026 17:02

mnorris11 mentioned this pull request Aug 6, 2026

[REVIEW] Add device-memory overload for cagra::helpers::optimize NVIDIA/cuvs#2423

Open

meta-codesync Bot changed the title ~~Fold multi-GPU CAGRA build into train(), delete trainMultiGpu~~ Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500) Aug 6, 2026

mnorris11 force-pushed the export-D114685755 branch from 558eb9c to 8da5b57 Compare August 6, 2026 18:26

mnorris11 force-pushed the export-D114685755 branch from 8da5b57 to 4e55b86 Compare August 6, 2026 18:35

mnorris11 force-pushed the export-D114685755 branch 2 times, most recently from 43ce419 to 4cf9428 Compare August 11, 2026 16:16

mnorris11 added the ignore label Aug 11, 2026

mnorris11 force-pushed the export-D114685755 branch from 4cf9428 to 931915a Compare August 12, 2026 15:32

mnorris11 force-pushed the export-D114685755 branch from 931915a to 95e2493 Compare August 13, 2026 04:05

mnorris11 force-pushed the export-D114685755 branch from 95e2493 to 7b4287f Compare August 13, 2026 15:40

mnorris11 force-pushed the export-D114685755 branch from 7b4287f to c92411c Compare August 13, 2026 20:41

mnorris11 force-pushed the export-D114685755 branch from c92411c to cd22982 Compare August 14, 2026 03:46

mnorris11 force-pushed the export-D114685755 branch from cd22982 to 7de34ee Compare August 14, 2026 14:42

mnorris11 force-pushed the export-D114685755 branch from 7de34ee to 5ac8360 Compare August 14, 2026 15:44

mnorris11 force-pushed the export-D114685755 branch from 5ac8360 to 9f4ea20 Compare August 14, 2026 15:47

meta-codesync Bot closed this in 6644dcf Aug 14, 2026

meta-codesync Bot added the Merged label Aug 14, 2026

facebook-github-tools Bot added the Reverted label Aug 15, 2026

Navigation Menu

Uh oh!

Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500) - #5500

Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500)#5500
mnorris11 wants to merge 1 commit into
facebookresearch:mainfrom
mnorris11:export-D114685755

mnorris11 commented Aug 4, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

meta-codesync Bot commented Aug 4, 2026

Uh oh!

meta-codesync Bot commented Aug 14, 2026

Uh oh!

facebook-github-tools Bot commented Aug 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

mnorris11 commented Aug 4, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented Aug 4, 2026

Uh oh!

meta-codesync Bot commented Aug 14, 2026

Uh oh!

facebook-github-tools Bot commented Aug 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

mnorris11 commented Aug 4, 2026 •

edited by meta-codesync Bot

Loading