Repository navigation
fix: Bound embedding input length and attention memory per batch - #837
Draft
buildingjoshbetter wants to merge 29 commits into
Draft
buildingjoshbetter wants to merge 29 commits into
buildingjoshbetter wants to merge 29 commits into
Conversation
# Conflicts: # tests/test-model-daemon-admission-297.py
# Conflicts: # tests/test-model-daemon-admission-297.py # truememory/embedding_target.py
The token budgets bounded attention by system memory only. On MPS the process cap is far smaller, so a batch with one long memory went MPS OOM and moved the model to CPU for the server lifetime. Batches on MPS now also fit the free share of the cap, using the measured peak footprint (48 bytes per batch x tokens^2 x head). Measured on Apple silicon at the 3.25 GiB cap with both models resident: 31 short texts plus one 860-token text and 32 texts of about 540 tokens both completed on MPS (4.0 s and 15.4 s; previously MPS OOM), cosine similarity >= 0.9996 against float32 CPU vectors. Reranking 100 candidates including ten at the 1,024-token cap completed on MPS in 6.7 s (previously MPS OOM).
# Conflicts: # tests/test-model-daemon-admission-297.py
…e waits The MPS allocation sizing was only exercised through the local encode helper. Pin it in the model server's embed and rerank loops, the head lookup through sentence-transformers modules, and the bookkeeping wait under a deadline shorter than its 0.25 s bound.
# Conflicts: # tests/test-model-daemon-admission-297.py # tests/test-model-server-memory-pressure.py
# Conflicts: # tests/test-model-daemon-admission-297.py
# Conflicts: # tests/test-model-daemon-admission-297.py
After an MPS OOM the server retried the failed piece, and every piece after it, at the size chosen for MPS. In the soak the first CPU piece was 11 texts of 1,718 tokens: about 4.16 GB of eager attention per layer, on top of 3.56 GB resident under a 4 GiB budget. CPU batches were otherwise sized by free system memory, which the server's budget does not track. accelerator_token_budget now sizes CPU batches so one batch's attention (8 bytes per score at the softmax peak, per head) fits 80% of the process RSS budget minus resident memory: the model server's owner limit, or the automatic budget in local processes. After an embed or rerank MPS OOM the server retries the failed piece and re-slices every remaining piece under that budget, keeping result order. Models whose head count is unknown keep their previous budget.
# Conflicts: # truememory/model_server.py
…xes-837-v1 # Conflicts: # truememory/embedding_target.py # truememory/memory_budget.py
test_cpu_models_ignore_the_mps_allocation checks that a CPU model's attention budget ignores the MPS allocation. The CPU budget is sized against the process RSS limit minus resident memory, so on a hosted runner after thousands of tests it came out lower (20,661,350 vs 33,554,432). The test now pins process_rss_limit high; the assertion is unchanged.
This branch has not been deployed
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.
Embedding models were run with their full input limit (32,768 tokens for Qwen3-Embedding-0.6B) and a fixed 32-text batch. Attention memory grows with batch × length², so one batch of long texts could need tens of GB. On a Mac it fails on MPS first and is then retried on CPU. This applies the same dynamic bound as the reranker change (#836) to every embedding path and tier.
encodecall, with no measurement, for single texts, tensor or list outputs, and batches that fit even at the smallest budget.Real Qwen3-Embedding-0.6B on CPU, 32 synthetic texts per call, 16 GiB container, peak cgroup memory before → after:
Single-query latency: 30.77 ms → 30.87 ms median over 200 calls (within noise).
Validation at aa37923, stacked on #836:
tests/test-model-daemon-admission-297.pypins the exact bytes ofmps_utils.pyto guard against unintended edits. The pin is updated to the reviewed bytes; themps_utils.pychange is the three guardedbudgeted_encodecall sites.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 aa37923:
embedding_target.pykeeps this branch's float32 and length caps with feat: Account for model memory and retain worker launch ownership #835's model repository lookup;memory_budget.pykeeps both new registry methods.Validation: whole suite on 8d02ecd (macOS, Python 3.14.6, SQLite 3.53.3, clean archive, isolated HOME): 4,741 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 (86bdf52 passes both files: 81 passed); the other 2 are the pre-existing NaN tests
test_issue_465_flag_set_immediatelyandtest_issue_485_nan_flag_set_after_completion.Base is
maintenance/reranker-token-budget-v1(#836). Draft; no merges or releases.