fix(cache): clone embeddings before save to avoid persisting whole ba…#13
Open
JiaqiL10 wants to merge 1 commit into
Open
fix(cache): clone embeddings before save to avoid persisting whole ba…#13JiaqiL10 wants to merge 1 commit into
JiaqiL10 wants to merge 1 commit into
Conversation
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.
Fixes #12
Compatibility / Impact
Non-breaking. Existing caches still load correctly (values are unchanged); no migration is required.
After the fix, newly built caches shrink from batch_size × dim to dim per file (e.g. 16× smaller at --batch-size 16, 128× at --batch-size 128). Old caches can optionally be rebuilt to reclaim disk.
Embedding values are bit-for-bit identical — no effect on any already-trained or in-progress models.
Train-time side benefit: torch.load no longer pulls the entire [batch, dim] storage into memory per sample, reducing DataLoader memory and disk I/O.
Verification
Build a cache before and after the fix from the same input and assert:
Values unchanged — loaded vectors are bitwise-equal.
Size fixed — per-file size is on the order of dim × 4 bytes and no longer scales with batch_size.
import torch, os
emb = torch.load(path, map_location="cpu")["all_findings"] # dict value
1) values: identical to a freshly-cloned compute (no change in numbers)
assert torch.equal(emb, emb.clone())
2) size: file no longer carries the whole batch storage
real = emb.numel() * emb.element_size() # dim * 4B (e.g. 16 KB)
print("file:", os.path.getsize(path), " real tensor:", real)
print("storage elems:", emb.untyped_storage().nbytes() // emb.element_size()) # == dim, not batchdim
print("storage_offset:", emb.storage_offset()) # 0 after fix (was idim)
Before the fix: file ≈ batch_size × dim × 4B, storage elems == batch_size × dim, storage_offset == i × dim.
After the fix: file ≈ dim × 4B, storage elems == dim, storage_offset == 0.