Skip to content

fix: Bound reranker input length and attention memory per batch - #836

Draft
buildingjoshbetter wants to merge 28 commits into
maintenance/remaining-code-only-v2from
maintenance/reranker-token-budget-v1
Draft

buildingjoshbetter wants to merge 28 commits into
maintenance/remaining-code-only-v2from
maintenance/reranker-token-budget-v1

Conversation

@buildingjoshbetter

@buildingjoshbetter buildingjoshbetter commented Oct 14, 2026 •

Copy link
Copy Markdown
Owner

The local reranker and the model server ran the cross-encoder with its full 8,192-token limit and a fixed 64-pair batch. Attention scratch grows with batch × length², so one request with long pairs could need tens of GB. On a Mac, that request first fails with an MPS out-of-memory and is then retried on CPU, where it keeps allocating.

This change:

  • caps reranker input at 1,024 tokens (truncated longest side first);
  • scores pairs longest-first in batches sized so batch × tokens² stays within the cost of 64 pairs of 512 tokens;
  • halves that budget below 6 GiB of available memory and quarters it below 2 GiB;
  • leaves the original single predict call untouched when every default batch already fits, and for models without a tokenizer and an integer max_seq_length (proxies, older or custom encoders).

Real ModernBERT reranker on CPU, 100 synthetic candidates, 16 GiB container, peak cgroup memory before → after:

Candidate length Before After Scores
50 words 0.86 GB 0.69 GB identical
400 words 1.39 GB 1.40 GB identical
1,500 words 3.87 GB 1.10 GB change (1,024-token cap)
6,000 words killed at 16 GiB 1.10 GB before never completed

Validation at f1f7655, stacked on #835:

  • Combined suite: 4,515 tests / 5,048 subtests passed (14 new tests), 19 skips, 3 network cases deselected.
  • Focused suites: 782 / 2,336 and 1,327 / 1,469 passed.
  • Installed wheel smoke: 94 Python source files verified.
  • Separate original gates: 739 native source/inert cases and 276 memory cases passed.
  • Independent review accepted after one reported defect was fixed.

Limitation: a cross-encoder that reports max_seq_length as None is neither capped nor budgeted. The shipped ModernBERT and MiniLM rerankers report integer limits.

Update: carries #835's audit fixes. This branch merges the current #835 head, so it includes every change listed there (model server memory recovery, SQLite safety, maintenance and foreground write fixes, close and backfill lifecycle, ingestion, tier switching, and the Haiku default for claude CLI extraction). See #835 for details and known limitations.

Own changes since f1f7655:

  • 83b50d9: rerankers that expose only max_length are capped as well.

Validation: whole suite on 6d45d15 (macOS, Python 3.14.6, SQLite 3.53.3, clean archive, isolated HOME): 4,716 passed, 58 skipped, 3 network cases deselected, 4,184 subtests. Two failures were tests still pinning the pre-M-v7-1 clusters identity, updated in #835 and merged here (899d4b9 passes both files: 81 passed); the other 2 are the pre-existing NaN tests test_issue_465_flag_set_immediately and test_issue_485_nan_flag_set_after_completion.

Base is maintenance/remaining-code-only-v2 (#835). Draft; no merges or releases.

# Conflicts:
#	truememory/model_server.py
CrossEncoder gained max_seq_length in sentence-transformers 5.5. Earlier
releases, still allowed by the dependency range, expose max_length (None
in 3.x until set), so the 1,024-token cap and the token budget were
skipped there. Read and cap whichever limit the model exposes.

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

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant