This repository was archived by the owner on Sep 30, 2024. It is now read-only.
Embeddings: optimize memory usage during decode#54972
Merged
Merged
Conversation
camdencheek
commented
Jul 14, 2023
| return nil, err | ||
| } | ||
|
|
||
| ei.Embeddings = make([]int8, 0, numChunks*ei.ColumnDimension) |
Member
Author
There was a problem hiding this comment.
First change: this we should be using embeddingsChunkSize instead of ColumnDimension because each chunk is not a vector, but rather a chunk of 10_000 floats.
camdencheek
commented
Jul 14, 2023
|
|
||
| ei.Embeddings = make([]int8, 0, numChunks*ei.ColumnDimension) | ||
| ei.Embeddings = make([]int8, 0, numChunks*embeddingsChunkSize) | ||
| embeddingsBuf := make([]float32, 0, embeddingsChunkSize) |
Member
Author
There was a problem hiding this comment.
Second change: move the decode buffer outside the loop so we don't allocate on every iteration.
camdencheek
commented
Jul 14, 2023
| return nil, err | ||
| } | ||
| ei.Embeddings = append(ei.Embeddings, Quantize(embeddingSlice)...) | ||
| ei.Embeddings = append(ei.Embeddings, Quantize(embeddingsBuf, quantizeBuf)...) |
Member
Author
There was a problem hiding this comment.
Third change: update Quantize to accept an optional buffer, which it will use if it's large enough to fit the output.
camdencheek
commented
Jul 14, 2023
| } | ||
| } | ||
|
|
||
| func BenchmarkCustomRepoEmbeddingIndexDownload(b *testing.B) { |
Member
Author
There was a problem hiding this comment.
I added this benchmark to prove to myself that this cuts down the allocations. See PR description for results.
gl-srgr
approved these changes
Jul 17, 2023
keegancsmith
approved these changes
Jul 17, 2023
keegancsmith
left a comment
Member
There was a problem hiding this comment.
I guess on the switch to a third party vector store we won't keep this store around? IE it isn't worth making improvements to how we store stuff?
Member
Author
|
Yeah, there is a lot more room for improvement here, but hopefully (🤞) this won't be around too much longer |
Contributor
|
The backport to To backport manually, run these commands in your terminal: # Fetch latest updates from GitHub
git fetch
# Create a new working tree
git worktree add .worktrees/backport-5.1 5.1
# Navigate to the new working tree
cd .worktrees/backport-5.1
# Create a new branch
git switch --create backport-54972-to-5.1
# Cherry-pick the merged commit of this pull request and resolve the conflicts
git cherry-pick -x --mainline 1 2880b4ffa0b8591d7a46ace259c8f522931dd617
# Push it to GitHub
git push --set-upstream origin backport-54972-to-5.1
# Go back to the original working tree
cd ../..
# Delete the working tree
git worktree remove .worktrees/backport-5.1Then, create a pull request where the |
camdencheek
added a commit
that referenced
this pull request
Jul 17, 2023
This makes memory usage ~constant during the decode process. Previously, we were undersizing the large allocation (which caused it to reallocate ~10 extra times) for the index and making a bunch of additional smaller allocations per loop. This meant we were spiking memory usage during decoding, and even if the garbage collector could keep up (which I wouldn't expect it to), we'd still need 2x the index size at the peak for live memory. This was causing an OOM. This PR accounts for a 2.2x reduction in allocated bytes and a 1.4x reduction in allocation count. This is quite significant when our indexes are many GBs in memory.
camdencheek
added a commit
that referenced
this pull request
Jul 17, 2023
This makes memory usage ~constant during the decode process. Previously, we were undersizing the large allocation (which caused it to reallocate ~10 extra times) for the index and making a bunch of additional smaller allocations per loop. This meant we were spiking memory usage during decoding, and even if the garbage collector could keep up (which I wouldn't expect it to), we'd still need 2x the index size at the peak for live memory. This was causing an OOM. This PR accounts for a 2.2x reduction in allocated bytes and a 1.4x reduction in allocation count. This is quite significant when our indexes are many GBs in memory.
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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
3 participants
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.
This makes memory usage ~constant during the decode process. Previously, we were undersizing the large allocation (which caused it to reallocate ~10 extra times) for the index and making a bunch of additional smaller allocations per loop. This meant we were spiking memory usage during decoding, and even if the garbage collector could keep up (which I wouldn't expect it to), we'd still need 2x the index size at the peak for live memory. This was causing an OOM.
Before:
After:
That's a 2.2x reduction in allocated bytes and a 1.4x reduction in allocation count. This is quite significant when our indexes are many GBs in memory.
Test plan
A new benchmark to demonstrate better memory behavior, a new quick test on the quantize changes, existing tests on encode/decode, and manual E2E tests for generating and searching with an embeddings index.