Skip to content

fix(cache): clone embeddings before save to avoid persisting whole ba…#13

Open
JiaqiL10 wants to merge 1 commit into
YalaLab:mainfrom
JiaqiL10:fix/clone-embedding-before-save
Open

fix(cache): clone embeddings before save to avoid persisting whole ba…#13
JiaqiL10 wants to merge 1 commit into
YalaLab:mainfrom
JiaqiL10:fix/clone-embedding-before-save

Conversation

@JiaqiL10

@JiaqiL10 JiaqiL10 commented Jun 4, 2026

Copy link
Copy Markdown

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 i
dim)
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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CachedTextEmbedding._save_embedding persists whole batch storage per .pt (≈batch_size× disk/RAM bloat)

1 participant