CAGRA API update and allow async host refinement #131

mfoerste4 · 2024-05-17T16:20:33Z

This PR updates the CAGRA public API, changes defaults, and improves refinement during IVF-PQ build step.

Updated defaults:

By default CAGRA would select NN descent. We fall back to IVF-PQ build algorithm if there is not enough memory for NN descent.
For the IVF-PQ build algo, the search params were updated to use n_probe = 0.01*nlist, and both LUT and internal distance type is set to fp16 (as opposed to fp8/fp32 previously)
By default build would create the index that contains both the graph and the dataset on GPU. If the dataset does not fit gpu, then the returned index will only contain the graph (on device). In such case the user is expected to call index.update_dataset() (for example with dataset in managed memory) before we can search the index.

API changes:

We can specify IVF-PQ build algo parameters the following way:

cagra::index_params params;
params.graph_degree = 32;
params.intermediate_graph_degree = 48;
auto pq_params = cagra::graph_build_algo::ivf_pq_params(dataset.extents());
// This sets reasonable defaults, but can be updated by the user, e.g.:
pq_params.pq_dim = 32;
// Select IVF-PQ algorithm by passing ivf_pq_params as graph_build_algo
params.graph_build_algo = pq_params;

We can specify NN descent algo parameters a similar way

cagra::index_params params;
params.graph_degree = 32;
params.intermediate_graph_degree = 48;
// Select NN descent algo by passing nn_descent_params
params.graph_build_algo = cagra::graph_build_params::nn_descent_params(intermediate_degree) ;

Additionally, this PR optimizes the IVF-PQ refinement step within the CAGRA graph creation in case the dataset is in host-memory.

depending on hardware, a certain amount of refinement comes for free
the change also has a positive effect without refinement as the graph creation is done on host either way

tfeher

Thanks Malte for the PR, it looks good overall, just two small comments.

Public API is needed to allow fine control of CAGRA IVF-PQ build options.

cpp/src/neighbors/detail/cagra/cagra_build.cuh

tfeher

Thanks Malte, for the updates! Let's discuss API questions with @cjnolet, otherwise, the refinement improvement part looks good!

tfeher · 2024-05-23T08:55:34Z

cpp/include/cuvs/neighbors/cagra.hpp

-           const cuvs::neighbors::cagra::index_params& params,
-           raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset)
-  -> cuvs::neighbors::cagra::index<uint8_t, uint32_t>;
+auto build(


I believe the documentation from cagra.cuh shall be moved here. @cjnolet to confirm.

Correct- we should document every function / prototype exposed in include/cuvs/**.hpp. It seems repetitive but the API docs are fully indexed and searchable so when a user finds one of these functions through search, it's guaranteed to have the docs carried along with it.

cpp/include/cuvs/neighbors/cagra.hpp

…archs

cjnolet · 2024-05-24T21:57:02Z

cpp/include/cuvs/neighbors/cagra.hpp

+  raft::resources const& handle,
+  const cuvs::neighbors::cagra::index_params& params,
+  raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
+  std::optional<cuvs::neighbors::nn_descent::index_params> nn_descent_params = std::nullopt,


This is awkward- a user is going to have to set std::nullopt for at least one of nn-descent_params or ivf_pq params. I think we should make these build_params for cagra. Actually, I think what we should do is take away the graph_build_algo option from build_params and add a graph_build_params option that accepts the cuvs::neighbors::build_params base class, documents it approprirately with what's allowed, and then dispatches accordingly. Also remember to raise an exception if anything other than ivf_pq or nn_descent is used. This allows us to add more graph build algos in the future without too much additional effort but it also makes the public APIs super simple for the user- by default we can just set the build_params to the default nn_descent::build_params.

On second though- since you need refine_ratio on the ivf_pq_build_params, why not have a new type cagra::graph_build_params which can extend cuvs::neighbors::build_params? Then you can have a cagra::graph_build_params::ivf_pq which extends ivf_pq::build_params and cagra::graph_build_params and adds the refine ratio. Then you can have cagra__graph_build_params::nn_descent which does the same for nn-descent. This way you can have cagra::build_params.graph_build_params which accepts anything that extends cagra::graph_build_params and dispatches the build algo accordingly? This would get all this additional stuff out of cagra::build and keep the API clean and easy to follow, while minimizing duplication.

@cjnolet , i am not sure whether I understand correctly what you are proposing. An abstract graph_build_params struct that is inherited from different variants, e.g.

namespace ivf_pq { struct graph_build_params : cuvs::neighbors::cagra::index_params { cuvs::neighbors::ivf_pq::index_params pq_build_params; cuvs::neighbors::ivf_pq::search_params pq_search_params; float refinement_rate = 2f; }; } // namespace

for the ivf_pq variant? To my understanding this would require to either implement interfaces individually or use of polymorphism/pointers that can be casted/dispatched. Do I understand your proposal correctly?

As an alternative cleanup step we could also combine additional variant-specific-parameters into (ivf_pq::)graph_build_params and add them as a single argument to the cagra::index_params

struct index_params : cuvs::neighbors::index_params { size_t intermediate_graph_degree = 128; size_t graph_degree = 64; std::optional<vpq_params> compression = std::nullopt; bool construct_index_with_dataset = true; std::variant<ivf_pq::graph_build_params, nn_descent::graph_build_params> graph_build_params; };

What do you think should we go for?

I think your last idea looks great. It's clean, explicit, and tells the user right away what graph_build_params variants are supported.

cjnolet · 2024-05-24T22:06:57Z

cpp/include/cuvs/neighbors/cagra.hpp

+  raft::resources const& handle,
+  const cuvs::neighbors::cagra::index_params& params,
+  raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
+  std::optional<cuvs::neighbors::nn_descent::index_params> nn_descent_params = std::nullopt,


On second though- since you need refine_ratio on the ivf_pq_build_params, why not have a new type cagra::graph_build_params which can extend cuvs::neighbors::build_params? Then you can have a cagra::graph_build_params::ivf_pq which extends ivf_pq::build_params and cagra::graph_build_params and adds the refine ratio. Then you can have cagra__graph_build_params::nn_descent which does the same for nn-descent. This way you can have cagra::build_params.graph_build_params which accepts anything that extends cagra::graph_build_params and dispatches the build algo accordingly? This would get all this additional stuff out of cagra::build and keep the API clean and easy to follow, while minimizing duplication.

cpp/src/neighbors/cagra.cuh

cpp/include/cuvs/neighbors/cagra.hpp

tfeher · 2024-05-29T10:59:25Z

cpp/include/cuvs/neighbors/cagra.hpp

  /**
   * Specify compression parameters if compression is desired.
   *
   * NOTE: this is experimental new API, consider it unsafe.
   */
  std::optional<cuvs::neighbors::vpq_params> compression = std::nullopt;
+


Suggested change

bool construct_index_with_dataset = true;

Removed, will try to put dataset on device. In case of OOM, catch error and return index without added dataset.

cpp/include/cuvs/neighbors/cagra.hpp

mfoerste4 · 2024-05-29T12:51:21Z

@tfeher , thanks for the review. I pushed a new version containing the suggestions.

tfeher

Thanks Malte, just one nitpick.

cpp/include/cuvs/neighbors/cagra.hpp

mfoerste4 · 2024-05-29T17:10:10Z

@cnolet, @tfeher, I just pushed a fix for the docs which was the last red CI check.

cjnolet · 2024-05-29T21:32:03Z

cpp/include/cuvs/neighbors/cagra.hpp

  /**
   * Specify compression parameters if compression is desired.
   *
   * NOTE: this is experimental new API, consider it unsafe.
   */
  std::optional<cuvs::neighbors::vpq_params> compression = std::nullopt;
+
+  bool construct_index_with_dataset = true;


cc @divyegala this is the option to enable not having the dataset vectors copied to device automatically after building the index.

Now this is set automatically: https://github.com/rapidsai/cuvs/pull/131/files#diff-f3ca3714f8c6d118b61d369e6edc085a4ff20ef17a5ec804beb53432e58f9940R479-R484

cpp/include/cuvs/neighbors/cagra.hpp

cjnolet · 2024-05-29T21:36:47Z

cpp/src/neighbors/cagra_c.cpp

+      index_params.graph_build_params = cuvs::neighbors::cagra::graph_build_params::ivf_pq_params{};
+      break;
+    case cuvsCagraGraphBuildAlgo::NN_DESCENT:
+      cuvs::neighbors::cagra::graph_build_params::nn_descent_params nn_descent_params{};


Can we consolidate the C++ and C structs so that we don't need to duplicate all of these values, please? It's going to be painful to manage two complete copies of these structs because changes are going to get made to one without being made to the other.

I did this with the DIstanceType enum in cuvs::distance, for example. I think we should ultimately strive to write these once (in C++ or C) and then expose them through the other.

tfeher

Highlighting recent changes.

tfeher · 2024-05-30T13:28:20Z

cpp/include/cuvs/neighbors/cagra.hpp

  /**
   * Specify compression parameters if compression is desired.
   *
   * NOTE: this is experimental new API, consider it unsafe.
   */
  std::optional<cuvs::neighbors::vpq_params> compression = std::nullopt;
+


Removed, will try to put dataset on device. In case of OOM, catch error and return index without added dataset.

tfeher · 2024-05-30T13:30:27Z

cpp/src/neighbors/detail/cagra/cagra_build.cpp

+  search_params.n_probes                = std::max<uint32_t>(10, build_params.n_lists * 0.01);
+  search_params.lut_dtype               = CUDA_R_16F;
+  search_params.internal_distance_dtype = CUDA_R_16F;


Updated defaults.

tfeher · 2024-05-30T13:31:46Z

cpp/src/neighbors/detail/cagra/cagra_build.cuh

-    build_knn_graph(res, dataset, knn_graph->view(), refine_rate, pq_build_params, search_params);
+  // Set default value in case knn_build_params is not defined.
+  auto knn_build_params = params.graph_build_params;
+  if (std::holds_alternative<std::monostate>(params.graph_build_params)) {


Select default algorithm

tfeher · 2024-05-30T13:32:33Z

cpp/src/neighbors/detail/cagra/cagra_build.cuh

-    // We just add the graph. User is expected to update dataset separately. This branch is used
-    // if user needs special control of memory allocations for the dataset.
+  } catch (std::bad_alloc& e) {
+    RAFT_LOG_DEBUG("Insufficient GPU memory to construct CAGRA index with dataset on GPU");


Automatically construct index without dataset if dataset would not fit GPU memory.

…ndex_async_refine

cjnolet · 2024-05-30T14:38:57Z

cpp/include/cuvs/neighbors/nn_descent.hpp

+ * @return true if enough GPU memory can be allocated
+ * @return false otherwise
+ */
+bool has_enough_device_memory(raft::resources const& res,


Small nitpick- can we move this to the cuvs::neighbors::nn_descent::helpers namespace please?

We can do it in a follow-on PR for 24.08 so we can get this merged ASAP.

cjnolet · 2024-05-30T15:40:08Z

/merge

Fixes handling OOM error during CAGRA index creation, that was introduced in #131. Authors: - Tamas Bela Feher (https://github.com/tfeher) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #167

This PR updates the CAGRA public API, changes defaults, and improves refinement during IVF-PQ build step. Updated defaults: - By default CAGRA would select NN descent. We fall back to IVF-PQ build algorithm if there is not enough memory for NN descent. - For the IVF-PQ build algo, the search params were updated to use `n_probe = 0.01*nlist`, and both LUT and internal distance type is set to `fp16` (as opposed to `fp8/fp32` previously) - By default build would create the index that contains both the `graph` and the `dataset` on GPU. If the `dataset` does not fit gpu, then the returned index will only contain the graph (on device). In such case the user is expected to call `index.update_dataset()` (for example with dataset in managed memory) before we can `search` the index. API changes: - We can specify IVF-PQ build algo parameters the following way: ```c++ cagra::index_params params; params.graph_degree = 32; params.intermediate_graph_degree = 48; auto pq_params = cagra::graph_build_algo::ivf_pq_params(dataset.extents()); // This sets reasonable defaults, but can be updated by the user, e.g.: pq_params.pq_dim = 32; // Select IVF-PQ algorithm by passing ivf_pq_params as graph_build_algo params.graph_build_algo = pq_params; ``` - We can specify NN descent algo parameters a similar way ```c++ cagra::index_params params; params.graph_degree = 32; params.intermediate_graph_degree = 48; // Select NN descent algo by passing nn_descent_params params.graph_build_algo = cagra::graph_build_params::nn_descent_params(intermediate_degree) ; ``` Additionally, this PR optimizes the IVF-PQ refinement step within the CAGRA graph creation in case the dataset is in host-memory. * depending on hardware, a certain amount of refinement comes for free * the change also has a positive effect without refinement as the graph creation is done on host either way Authors: - Malte Förster (https://github.com/mfoerste4) - Tamas Bela Feher (https://github.com/tfeher) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Tamas Bela Feher (https://github.com/tfeher) URL: rapidsai/cuvs#131

mfoerste4 added 3 commits May 17, 2024 12:35

enable asynchronous host-refinement for cagra index build

4416b0b

move alloc out of OMP region

946e106

added test for refinement -- disabled until detail::build API is exposed

7c6644e

github-actions bot added the cpp label May 17, 2024

cjnolet assigned mfoerste4 May 17, 2024

cjnolet added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels May 17, 2024

tfeher requested changes May 22, 2024

View reviewed changes

cpp/src/neighbors/detail/cagra/cagra_build.cuh Outdated Show resolved Hide resolved

cpp/src/neighbors/detail/cagra/cagra_build.cuh Outdated Show resolved Hide resolved

mfoerste4 added 3 commits May 22, 2024 21:22

expose more cagra build parameters

2c3f5e4

clarify code

4d5ff0c

fix merge conflict

545032b

mfoerste4 changed the title ~~[WIP] CAGRA/IVF-PQ index allow async host refinement~~ CAGRA/IVF-PQ index allow async host refinement May 22, 2024

mfoerste4 requested a review from tfeher May 22, 2024 22:21

mfoerste4 marked this pull request as ready for review May 23, 2024 08:40

mfoerste4 requested a review from a team as a code owner May 23, 2024 08:40

tfeher reviewed May 23, 2024

View reviewed changes

slightly reduce expected recall to catch refinement 1 results on all …

d2c1dad

…archs

cjnolet reviewed May 24, 2024

View reviewed changes

cjnolet requested changes May 24, 2024

View reviewed changes

tfeher mentioned this pull request May 28, 2024

ANN_BENCH #130

Merged

10 tasks

mfoerste4 added 5 commits May 29, 2024 08:30

cleanup API - use std::variant to switch algo of cagra graph build

9823831

fix c-api nn_descent

f310971

move documentation to API header

c1136c5

resolve merge conflicts

bb87c97

correct typo

8506c6a

tfeher requested changes May 29, 2024

View reviewed changes

review suggestions

3a6e776

mfoerste4 requested a review from tfeher May 29, 2024 12:48

tfeher requested changes May 29, 2024

View reviewed changes

cpp/include/cuvs/neighbors/cagra.hpp Outdated Show resolved Hide resolved

cpp/include/cuvs/neighbors/cagra.hpp Outdated Show resolved Hide resolved

mfoerste4 requested a review from cjnolet May 29, 2024 17:08

cjnolet reviewed May 29, 2024

View reviewed changes

cpp/include/cuvs/neighbors/cagra.hpp Outdated Show resolved Hide resolved

cjnolet reviewed May 29, 2024

View reviewed changes

cpp/include/cuvs/neighbors/cagra.hpp Outdated Show resolved Hide resolved

cjnolet reviewed May 29, 2024

View reviewed changes

Update CAGRA API (NN descent build by default)

95f6064

tfeher requested a review from a team as a code owner May 30, 2024 12:59

github-actions bot added the CMake label May 30, 2024

tfeher changed the title ~~CAGRA/IVF-PQ index allow async host refinement~~ CAGRA API update and allow async host refinement May 30, 2024

tfeher reviewed May 30, 2024

View reviewed changes

Fix style

3e22b13

cjnolet approved these changes May 30, 2024

View reviewed changes

tfeher and others added 5 commits May 30, 2024 16:17

remove cagra's build_host and build_device functions

fd66402

remove debug printouts

8b1fbb8

Merge remote-tracking branch 'origin/branch-24.06' into cagra_ivfpq_i…

e688453

…ndex_async_refine

fix style

f5dff7d

Merge branch 'branch-24.06' into cagra_ivfpq_index_async_refine

71c779a

cjnolet reviewed May 30, 2024

View reviewed changes

cjnolet approved these changes May 30, 2024

View reviewed changes

cjnolet requested a review from tfeher May 30, 2024 14:39

Fix style

2540dbf

tfeher approved these changes May 30, 2024

View reviewed changes

Fix CAGRA C API default algo params

72d7b62

rapids-bot bot merged commit 13e0732 into rapidsai:branch-24.06 May 30, 2024
54 checks passed

tfeher mentioned this pull request May 30, 2024

Fix CAGRA OOM handling #167

Merged

tfeher mentioned this pull request Aug 12, 2024

[FEA] Add CAGRA graph build options to the C and Python API #289

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

CAGRA API update and allow async host refinement #131

CAGRA API update and allow async host refinement #131

mfoerste4 commented May 17, 2024 •

edited by tfeher

Loading

tfeher left a comment

tfeher left a comment

tfeher May 23, 2024

cjnolet May 24, 2024 •

edited

Loading

cjnolet May 24, 2024 •

edited

Loading

cjnolet May 24, 2024

mfoerste4 May 27, 2024

cjnolet May 29, 2024

cjnolet May 24, 2024

tfeher May 29, 2024

tfeher May 30, 2024

mfoerste4 commented May 29, 2024

tfeher left a comment

mfoerste4 commented May 29, 2024

cjnolet May 29, 2024

tfeher May 30, 2024

cjnolet May 29, 2024

tfeher left a comment

tfeher May 30, 2024

tfeher May 30, 2024

tfeher May 30, 2024

tfeher May 30, 2024

cjnolet May 30, 2024

cjnolet commented May 30, 2024

CAGRA API update and allow async host refinement #131

CAGRA API update and allow async host refinement #131

Conversation

mfoerste4 commented May 17, 2024 • edited by tfeher Loading

tfeher left a comment

Choose a reason for hiding this comment

tfeher left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cjnolet May 24, 2024 • edited Loading

Choose a reason for hiding this comment

cjnolet May 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mfoerste4 commented May 29, 2024

tfeher left a comment

Choose a reason for hiding this comment

mfoerste4 commented May 29, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tfeher left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cjnolet commented May 30, 2024

mfoerste4 commented May 17, 2024 •

edited by tfeher

Loading

cjnolet May 24, 2024 •

edited

Loading

cjnolet May 24, 2024 •

edited

Loading