studio: charge the tied-embedding output duplicate to the VRAM budget - #9929
studio: charge the tied-embedding output duplicate to the VRAM budget#9929danielhanchen wants to merge 2 commits into
Conversation
A model that ties its input and output embeddings ships no output.weight. llama.cpp does not reuse token_embd in place: it re-creates the output tensor from it as TENSOR_DUPLICATED, and a second vocabulary matrix is really allocated. The load therefore needs the file's tensors plus one more copy of the embedding matrix, and the context budget was sizing it from the GGUF file size alone. Under-counting weights is the dangerous direction. It leaves the context search believing there is VRAM that the load will consume, so Auto picks a context the card cannot hold. Measured on gemma-4-E2B-it UD-Q4_K_XL: 3021.88 MiB of tensors in the file, 3285.89 MiB of model buffers reported by llama-server. The difference is 264.01 MiB against a token_embd of exactly 264.00 MiB. The two copies land on different devices in that same run, the original in CPU_Mapped and the duplicate in CUDA0, which is why the duplicate is a VRAM cost rather than a RAM one, and why the split is worth making rather than aliasing: the input side is a row gather that is happy on the host, the output side is a vocabulary GEMM per token that wants the GPU. Across the shipped quants the charge is 264 MiB (E2B UD-Q4_K_XL) to 924 MiB (31B UD-Q4_K_XL), 4.6% to 8.7% of file size on the gemma family. Qwen3.6 and Qwen3.8 ship a real output.weight and are unaffected at 0 bytes. A split GGUF abstains: GGUFReader maps only the path it is given, so an output.weight living in a later shard would read as tied here and add a duplicate that is never allocated. Anything unreadable returns 0, so the budget keeps the old behaviour rather than failing a launch. The probe is cached on file identity, not path, because the context search calls it once per candidate context and a model replaced in place must not serve its predecessor's answer.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
π‘ Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e20e0732c
βΉοΈ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with π.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # card cannot hold. 264 MiB on gemma-4-E2B UD-Q4_K_XL and | ||
| # 924 MiB on gemma-4-31B UD-Q4_K_XL, against files of 3037 | ||
| # and 17951 MiB. | ||
| model_size = gguf_size + mmproj_size + self._tied_output_bytes(model_path) |
There was a problem hiding this comment.
Preserve the tied-output charge when pinning the projector
For a tied-embedding vision model that enters the automatic projector-placement path, this charge is not propagated through the rest of the budget: _mm_need at line 18697 is still built from bare gguf_size, and if the projector is pinned, _apply_mmproj_cpu_pin resets model_size to bare gguf_size at line 18617. Consequently, the exact Gemma-style duplicate added here disappears from the subsequent context search, which can again select a context that exceeds VRAM or incorrectly keep the projector on the GPU. Store the duplicate/base-model footprint separately and retain it in both calculations.
Useful? React with πΒ / π.
| if name == "output.weight": | ||
| return 0 | ||
| if name == "token_embd.weight": | ||
| embd = int(tensor.n_bytes) | ||
| return embd |
There was a problem hiding this comment.
Exclude encoder-only GGUFs from tied-output charging
Absence of output.weight does not by itself prove tied language-model embeddings. Studio explicitly supports encoder-only embedding GGUFs such as BERT and nomic-bert through is_embedding_gguf; these models have token_embd.weight but do not create a vocabulary-output tensor, so this loop charges an entire nonexistent duplicate. On those loads Auto can unnecessarily reduce context, spill layers, or reject tensor placement. Inspect general.architecture/pooling metadata and only apply this charge to architectures whose llama.cpp loader actually synthesizes a tied output.
Useful? React with πΒ / π.
| return LlamaCppBackend._tied_output_bytes_cached( | ||
| model_path, stat.st_size, stat.st_mtime_ns, | ||
| ) |
There was a problem hiding this comment.
Key the probe cache on actual file identity
The cache key described here as file identity contains only path, size, and mtime. If a GGUF is atomically replaced by a same-sized file with a preserved timestamp, or rewritten in place while restoring its mtime, the cached tied/untied result survives even though the tensor table changed; a stale zero or stale embedding size then produces an incorrect context budget for every subsequent load. The backend's _gguf_load_source_identity already uses device/inode in addition to size and mtime, so include inode/device and preferably ctime here as well.
Useful? React with πΒ / π.
|
Confirmed the budget in studio/backend/core/inference/llama_cpp.py still sizes a load as gguf_size + mmproj_size, so tied-embedding models lose a whole vocabulary matrix of VRAM and Auto over-picks context. Will get this reviewed. |
What
The context budget sizes a load's weights from the GGUF file size. For a model that ties its input and output embeddings, that under-counts by one vocabulary matrix.
Such a model ships no
output.weight. llama.cpp does not reusetoken_embdin place: it re-creates the output tensor from it asTENSOR_DUPLICATED, and a second vocabulary matrix is really allocated.Why it matters
Under-counting weights is the dangerous direction. It leaves the context search believing there is VRAM that the load will then consume, so Auto picks a context the card cannot hold.
The measurement
gemma-4-E2B-it UD-Q4_K_XL:
token_embd.weightThe two copies land on different devices in that same run.
CPU_Mappedis 1804.00 MiB, which is the per-layer embeddings (1540) plustoken_embd(264);CUDA0is 1481.89 MiB, of which the non-embedding tensors account for 1217.88, leaving 264.01. So the original stays host-resident and the duplicate sits in VRAM, which is exactly the quantity this budget is spending.That also makes the duplication reasonable rather than wasteful: the input side is a row gather that is happy on the host, while the output side is a vocabulary GEMM every token that wants to be on the GPU. Aliasing one tensor would force both roles onto whichever device won.
Across the shipped quants:
Qwen ships a real
output.weight, so it is unaffected.Behaviour
Auto context becomes slightly smaller on tied models, which is the correction. Nothing changes for models that ship their own output tensor.
Edge cases
GGUFReadermaps only the path it is given, so anoutput.weightliving in a later shard would read as tied and add a duplicate that is never allocated. That is an over-count, on exactly the models large enough to be split, so a shard returns 0.Tests
studio/backend/tests/test_tied_output_vram_budget.py, 7 tests, no GPU or network. They build real GGUFs in a tmp dir rather than mocking the reader, so the probe is exercised through the same path the product uses. Covered: a tied model is charged the whole matrix, an untied one is charged nothing, the charge scales with vocabulary size rather than being a constant, a shard abstains, junk and missing files return 0, an in-place swap is not served from cache, and the call site still adds it.409 tests pass across the surrounding budget suites (
test_auto_offload_ctx_*,test_mtp_vram_budget,test_vram_budget_settings) plus the new file.