Fix GPU and Java-heap OOM when building large vector indexes#173
Open
nvzm123 wants to merge 3 commits into
Open
Fix GPU and Java-heap OOM when building large vector indexes#173nvzm123 wants to merge 3 commits into
nvzm123 wants to merge 3 commits into
Conversation
…onversion from CAGRA to HNSW
without loading the full set of data on the Java Heap, but instead allows us to stream the set of data to the Java Heap.
Address review feedback on duplicated code paths introduced by the host-streaming refactor: - Merge the two createMultiLayerHnswGraph overloads (List-based and CuVSMatrix-based) into a single CuVSMatrix-based method. The flush and quantized paths now build a matrix and sample rows from it, matching the merge path. - Merge the two writeFieldInternal overloads into one that takes a CuVSMatrix, with a thin List adapter for the flush/sorting path. - Drop the redundant `size` parameter throughout; it is derived from dataset.size(). - Remove the now-unused CuVSResources parameter from the Utils matrix builders (createFloatMatrix/createByteMatrix/createByteMatrixFromArray) and correct their docs, which still referred to device memory after the switch to hostBuilder. Updated all call sites (CuVS2510GPUVectorsWriter, binary/scalar quantized writers). While migrating the quantized writers to the unified graph method, fixed a latent byte-width bug: the quantized branch hardcoded (dimensions + 7) / 8 (binary packing), which is wrong for scalar quantization (one byte per dimension). The width is now taken from the matrix's column count, correct for both binary and scalar. Signed-off-by: Zack Meeks <zmeeks@nvidia.com>
Author
|
I think the two main fixes you made are essential. Thanks, @nvzm123! Re: the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes two separate out-of-memory failures when building accelerated HNSW
indexes over large datasets:
GPU / device-memory OOM: the dataset matrix was built with
CuVSMatrix.deviceBuilder, which eagerly loads the entire dataset onto theGPU. This switches to
CuVSMatrix.hostBuilder, which keeps the dataset inhost memory and lets CAGRA stream it host→device as needed, rather than
requiring the full dataset resident on the GPU at once.
Java-heap OOM: HNSW graph construction previously materialized the full
dataset on the Java heap (
List<float[]>). It now streams subsets from anative host matrix during graph construction, so the full dataset is never
held on the heap.
Supersedes #141. This branch was cut from latest
mainfor a clean rebase andcontains all of #141's commits (through 7c5c448) plus the review-feedback
consolidation described below.
Review feedback from #141
Addresses @narangvivek10's four review comments:
Duplicate
createMultiLayerHnswGraph(List vsCuVSMatrixoverloads) —consolidated into a single
CuVSMatrix-based method inAcceleratedHNSWUtils. All callers (flush, merge, and the binary/scalarquantized writers) now build a matrix and pass it. Migrating the quantized
writers surfaced a latent byte-width bug: the quantized branch hardcoded
(dimensions + 7) / 8(binary packing), which is wrong for scalarquantization (one byte per dimension). Width is now taken from the matrix's
column count, correct for both binary and scalar.
Redundant
sizeparameter — removed fromwriteFieldInternalandcreateMultiLayerHnswGraph; derived internally fromdataset.size().Duplicate
writeFieldInternal— consolidated into a singlewriteFieldInternal(FieldInfo, CuVSMatrix)inLucene99AcceleratedHNSWVectorsWriter, with a thinListadapter for theflush/sorting path.
Utilshost/device flexibility — theCuVSResourcesparameter on thematrix builders was unused after the switch to
hostBuilder, so it'sremoved, and the docs that still referred to "device memory" are corrected.
Removing it updated call sites in
CuVS2510GPUVectorsWriterand thebinary/scalar quantized writers. Making the builders fully
host/device-selectable seems broader than this PR — device construction
still exists directly in
CuVS2510GPUVectorsReader, so a properhost/device-aware API would want to account for that. Would you prefer I
file a separate issue to track it, or leave it until something actually
needs device-side construction through these helpers? Happy either way.
Testing
All format and merge suites pass locally:
TestMerge(includingtestLargeScaleMergewith-DlargeScale=true)TestCuVSVectorsFormatTestLucene99AcceleratedHNSWVectorsFormatTestQuantizedVectorsFormatsFlaky test note:
TestLucene99AcceleratedHNSWVectorsFormat.testCheckIntegrityReadsAllBytesfails intermittently (~2 in 6 full-suite runs) via
CheckIndexintestRandomWithUpdatesAndGraph. It passes 100% in isolation and fails at thesame rate on the base branch without the consolidation, so it appears to be a
pre-existing, seed-dependent suite-level issue rather than something introduced
here. Flagging in case CI hits it.