Support graph-based approximate nearest neighbors search (cagra on nn_descent/ivf_pq) #714
Conversation
|
build |
lijinf2
changed the title
|
build |
python/src/spark_rapids_ml/knn.py
Outdated
| * intermediate_graph_degree: (int, default = 128) an intermediate variable used during graph index construction | ||
| * graph_degree: (int, default = 64) the degree of each node in the final graph index | ||
|
|
||
| parameters for search, full list in cuvs python API documentation. |
There was a problem hiding this comment.
Would be great to have a clickable link (in our api docs github pages site) to this. Let's make a note to add one in the future. Need to figure how to do it in Sphinx.
There was a problem hiding this comment.
Yeah, added the link and tested new docs locally.
python/src/spark_rapids_ml/knn.py
Outdated
| parameters for index construction | ||
| * build_algo: (str, default = 'ivf_pq') algorithm to build graph index, can be either 'ivf_pq' or 'nn_descent'. nn_descent is expected to be generally faster than ivf_pq. | ||
| * intermediate_graph_degree: (int, default = 128) an intermediate variable used during graph index construction | ||
| * graph_degree: (int, default = 64) the degree of each node in the final graph index |
There was a problem hiding this comment.
Is this related to the top level topk parameter? If so, should not be exposed (nor documented) separately and should be set in our code based on the top level parameter.
There was a problem hiding this comment.
It seems irrelevant that can be set independently, as I can understand.
graph_degree is the number of neighbors of every node in the final graph index. It is used in the index construction phase to control the size of graph index.
topk is used in the search phase to control the number of returned neighbors.
When a query visits a node of the graph index, all of its neighbors will be evaluated to update the current topk results. A larger graph_degree should be helpful for a given k, and the two parameters can be set independently.
python/src/spark_rapids_ml/knn.py
Outdated
| * graph_degree: (int, default = 64) the degree of each node in the final graph index | ||
|
|
||
| parameters for search, full list in cuvs python API documentation. | ||
| * itopk_size: (int, default = 64) number of intermediate search results retained during the search. Larger value improves the search accuracy but increases the search time. |
There was a problem hiding this comment.
I guess the intermediate topk related params need to be no smaller than the top level one? Our code should enforce/set this.
There was a problem hiding this comment.
Yes, added validation on the itopk_size parameter. cuml increases the value to the nearest multiple of 32 and requires the value to be larger than or equal to k.
|
build |
python/src/spark_rapids_ml/knn.py
Outdated
| internal_topk_size = math.ceil(itopk_size / 32) * 32 | ||
| if internal_topk_size < topk: | ||
| raise ValueError( | ||
| f"cuVS increases itopk_size to be closest multiple of 32 and expects the value, i.e. {internal_topk_size}, to be larger than or equal to k, i.e. {topk})." |
There was a problem hiding this comment.
is more correct to say cagra increase?
| n_neighbors = 50 | ||
| error_msg = "" | ||
| if internal_topk_size < n_neighbors: | ||
| error_msg = f"cuVS increases itopk_size to be closest multiple of 32 and expects the value, i.e. {internal_topk_size}, to be larger than or equal to k, i.e. {n_neighbors}." |
There was a problem hiding this comment.
same comment as above.
Signed-off-by: Jinfeng <[email protected]>
|
build |
No description provided.