Skip to content

lore: vectorize emails 64 at a time - #122

Open
vladwing wants to merge 1 commit into
facebookexperimental:mainfrom
vladwing:lore-vectors-chunked
Open

vladwing wants to merge 1 commit into
facebookexperimental:mainfrom
vladwing:lore-vectors-chunked

Conversation

@vladwing

Copy link
Copy Markdown
Contributor

On a production-sized lore database, tokenizing at most 64 emails per model call cuts cold --vectors peak memory from 32-36 GB to 22 GB on main, and from 22-27 GB to 10.4-10.6 GB with #112 and #113 applied, with the same vectors and no added wall time.

update_lore_vectors() starts one worker per CPU, and each worker passed its whole RecordBatch, up to 1,024 emails, to vectorize_batch() in one call. model2vec's chunk size defaults to num_cpus * 128, larger than any batch, so the tokenizer held the untruncated encoding of every email in the batch until it had pooled them all. A lance batch never spans fragments, so the peak follows fragment size. Main writes lore in small fragments, where the gain is modest (4.3-5.1 GB to 4.2 GB on two archives); #112's larger inserts, and any compaction of the lore table, write large ones, where 316 workers can hold the whole table at once. The 1,039- and 93-fragment databases measured below were built with earlier revisions of #112 and #113, the 632-fragment one by main.

Each worker now encodes its batch 64 emails at a time. model2vec pools every email on its own, and the tokenizer of the model that scripts/nomic2vec.py produces adds no batch padding, so the vectors do not depend on the chunking. A call that returns the wrong number of vectors is an error rather than a silent shift of every later vector onto the wrong message_id. The global SEMCODE_BATCH_SIZE default stays as it is, because it also sets the chunk size for code vectors; it still applies inside each lore call, so it can lower the lore chunk but not raise it.

Measured with semcode-index --vectors -j 32 on reflinked copies, on a shared 316-CPU host (1-minute load 66-227). "#112 + #113" is main merged with the PR heads 0fad419 and 4d66a45.

                                  peak RSS (GB)    wall (s)
production-sized: 995,564 emails, 1,039 fragments
  main                            35.85, 32.24     931, 967
  main + this commit              22.09, 22.09     932, 936
  main + #112 + #113              27.03, 22.21     167, 307
  main + #112 + #113 + this       10.38, 10.56     154, 256
two archives: 104,189 emails, 93 fragments
  main                            7.85, 8.09       25, 23
  main + this commit              3.30, 3.38       23, 23
two archives: 104,189 emails, 632 fragments
  main                            4.33, 5.09       38, 39
  main + this commit              4.20, 4.23       33, 36

The sorted (message_id, vector) digest of lore_vectors matches one call per batch on the production-sized and 93-fragment databases. Main's 930 s production wall time comes from elsewhere in the run: with #112 and #113 the same run takes 154-307 s. Of the gap between 22 GB on main and 10.4-10.6 GB with both PRs, about 6 GB is #112's cap of 32 tokio workers: #113 alone with this commit peaks at 16.7-17.1 GB, and at 10.8 GB with TOKIO_WORKER_THREADS=32.

On a production-sized lore database, tokenizing at most 64 emails per
model call cuts cold `--vectors` peak memory from 32-36 GB to 22 GB on
main, and from 22-27 GB to 10.4-10.6 GB with facebookexperimental#112 and facebookexperimental#113 applied, with
the same vectors and no added wall time.

update_lore_vectors() starts one worker per CPU, and each worker passed
its whole RecordBatch, up to 1,024 emails, to vectorize_batch() in one
call.  model2vec's chunk size defaults to num_cpus * 128, larger than
any batch, so the tokenizer held the untruncated encoding of every email
in the batch until it had pooled them all.  A lance batch never spans
fragments, so the peak follows fragment size.  Main writes lore in small
fragments, where the gain is modest (4.3-5.1 GB to 4.2 GB on two
archives); facebookexperimental#112's larger inserts, and any compaction of the lore table,
write large ones, where 316 workers can hold the whole table at once.
The 1,039- and 93-fragment databases measured below were built with
earlier revisions of facebookexperimental#112 and facebookexperimental#113, the 632-fragment one by main.

Each worker now encodes its batch 64 emails at a time.  model2vec pools
every email on its own, and the tokenizer of the model that
scripts/nomic2vec.py produces adds no batch padding, so the vectors do
not depend on the chunking.  A call that returns the wrong number of
vectors is an error rather than a silent shift of every later vector
onto the wrong message_id.  The global SEMCODE_BATCH_SIZE default stays
as it is, because it also sets the chunk size for code vectors; it still
applies inside each lore call, so it can lower the lore chunk but not
raise it.

Measured with `semcode-index --vectors -j 32` on reflinked copies, on a
shared 316-CPU host (1-minute load 66-227).  "facebookexperimental#112 + facebookexperimental#113" is main
merged with the PR heads 0fad419 and 4d66a45.

                                      peak RSS (GB)    wall (s)
    production-sized: 995,564 emails, 1,039 fragments
      main                            35.85, 32.24     931, 967
      main + this commit              22.09, 22.09     932, 936
      main + facebookexperimental#112 + facebookexperimental#113              27.03, 22.21     167, 307
      main + facebookexperimental#112 + facebookexperimental#113 + this       10.38, 10.56     154, 256
    two archives: 104,189 emails, 93 fragments
      main                            7.85, 8.09       25, 23
      main + this commit              3.30, 3.38       23, 23
    two archives: 104,189 emails, 632 fragments
      main                            4.33, 5.09       38, 39
      main + this commit              4.20, 4.23       33, 36

The sorted (message_id, vector) digest of lore_vectors matches one call
per batch on the production-sized and 93-fragment databases.  Main's
930 s production wall time comes from elsewhere in the run: with facebookexperimental#112
and facebookexperimental#113 the same run takes 154-307 s.  Of the gap between 22 GB on main
and 10.4-10.6 GB with both PRs, about 6 GB is facebookexperimental#112's cap of 32 tokio
workers: facebookexperimental#113 alone with this commit peaks at 16.7-17.1 GB, and at
10.8 GB with TOKIO_WORKER_THREADS=32.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Nov 11, 2026

This branch has not been deployed

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

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant