fix: key the built-model cache on device and inference dtype by jmkuebler · Pull Request #1219 · PriorLabs/TabPFN

jmkuebler · 2026-08-31T12:30:50Z

Issue

RES-2743. Stacked on #1250, which removes the dead cache_trainset_representation flag and lets fit_mode="fit_with_cache" use the cache.

Motivation and Context

The built-model cache is opt-in via TABPFN_MODEL_CACHE_SIZE and hands one built model to every estimator that asks for it, so a warm fit costs ~0.003 s instead of ~0.12 s.

It could not be shared across placements. An estimator applies the placement to the module it is handed, in place — the inference engine moves it with model.to(device) and casts it with model.type(dtype). Neither was part of the key, so two estimators wanting different placements were served the same instance:

a = TabPFNClassifier(device="cpu").fit(X, y); a.predict_proba(X_test)
b = TabPFNClassifier(device="mps").fit(X, y); b.predict_proba(X_test)
a.predict_proba(X_test)
# RuntimeError: Placeholder storage has not been allocated on MPS device!

TabPFNClassifier(inference_precision=torch.float16).fit(X, y).predict_proba(X_test)
TabPFNClassifier().fit(X, y).predict_proba(X_test)
# RuntimeError: mat1 and mat2 must have the same dtype, but got Float and Half

The dtype cannot be repaired instead of keyed: model.type(torch.float16) is destructive, so a model cast to fp16 can never serve a full-precision caller.

The suite is the clearest evidence

Setting TABPFN_MODEL_CACHE_SIZE over tests/test_classifier_interface.py and tests/test_regressor_interface.py:

	result	time
cache off	309 passed	146.4 s
cache on, before this change	38 failed	—
cache on, after	309 passed	116.8 s

Every one of the 38 was a placement collision — Float and Double, Float and Half, Cannot compare two tensors on different devices, Cannot convert a MPS Tensor to float64. The suite spans cpu/mps and three forced precisions, so it could not opt in at all before. On top of #1250, where fit_with_cache also hits the cache, the same suites show 64 collisions before this change and none after.

Those two files are unusually fit-dense. Over the whole suite the effect is smaller — 1861 passed either way, 398 s off versus 368 s on, so ~30 s or 7.5%. Enabling it in conftest.py is a follow-up, not part of this PR.

What a caller still owns

The default is unchanged, so nothing happens unless the env var is set. Mixing devices or precisions now yields separate entries rather than corruption. Two estimators served the same entry still share the module, so moving one with .to() after fit or training its weights reaches the other — deliberate, and left to the caller. Calling .to() before fit is safe: the key is built from the device the estimator has at fit time.

Public API Changes

No Public API changes
Yes, Public API changes (Details below)
tabpfn.model_loading.load_model and load_model_criterion_config take optional devices and force_inference_dtype keyword arguments (both default None, so existing calls are unaffected — a caller that places the returned model should pass them).
A non-integer TABPFN_MODEL_CACHE_SIZE warns instead of being silently ignored.
The cache default is unchanged (off).

How Has This Been Tested?

Full suite in the shipping configuration (cache off): 1861 passed, 176 skipped, 8 xfailed.
Full suite with TABPFN_MODEL_CACHE_SIZE=8: the same 1861 passed, 176 skipped, 8 xfailed, so the cache is exercised across every device and precision the suite covers.
After rebasing on Remove the dead cache_trainset_representation flag #1250: tests/test_classifier_interface.py, tests/test_regressor_interface.py and tests/test_model_cache.py with TABPFN_MODEL_CACHE_SIZE=8 on cpu and cuda (A100): all pass.
tests/test_model_cache.py: separate entries per device and per dtype; unspecified devices kept apart from device-specific ones; disabled by default; two entries holding a classifier and a regressor; invalid env value; LRU eviction; cache-clear forcing a rebuild; and a tripwire on load_model's parameters so a future parameter cannot quietly bypass the key.
Checked end-to-end on cpu, mps and cuda with the cache on: fitting on one device then the other leaves both estimators correct, and fp16 then full precision works.

Checklist

The changes have been tested locally.
Documentation has been updated (if the public API or usage changes).
A changelog entry has been added (see changelog/README.md), or "no changelog needed" label requested.
The code follows the project's style guidelines.
I have considered the impact of these changes on the public API.

🤖 Generated with Claude Code

jmkuebler · 2026-08-31T15:00:39Z

@cursor review

The built-model cache hands one instance to every estimator with the same placement, so anything that mutates a model in place reaches every other holder. Two such mutations were reachable now that the cache is on by default: - `.to()` on a fitted estimator moved the shared instance, so another estimator holding it failed its next predict with "Placeholder storage has not been allocated on MPS device!". Clearing the cache did not help: it stops future mis-placed hits but cannot un-share what was already handed out. - Fine-tuning optimises `model_`'s weights in place. A `FinetunedTabPFN*` fit followed by an ordinary fit in the same process silently returned the fine-tuned weights — no error, just wrong predictions. Both are the same rule: a caller about to mutate a model takes a private copy first. `detach_models_from_cache` does that for the models and the regressor's bar distribution, which comes from the same cached tuple, and is called from the four places that mutate: `TabPFNClassifier.to`, `TabPFNRegressor.to`, `fit_from_preprocessed` on both, and the finetuner before its optimizer binds to `model_`. Reported by Cursor Bugbot on #1219. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

`fit_from_preprocessed` only detached inside its re-initialisation branch, so any path reaching it with `models_` already populated — an earlier `fit()`, with the default `no_refit=True` — trained the shared cached instance. Moving the call out of the branch means it runs once per batch, which a plain deep copy could not survive: it would copy every step and hand back a different object than the optimizer is holding. So `detach_models_from_cache` now copies only while `is_cached_model` reports the cache still holds the instance, making it a no-op once the models are private and safe to call unconditionally. Reported by Cursor Bugbot on #1219. Its stated route was `get_preprocessed_datasets`, which no longer exists on the estimators (only stale docstring references remain), but the branch it identified was reachable by other means. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

jmkuebler · 2026-09-01T06:18:56Z

@cursor review

jmkuebler · 2026-09-01T09:22:43Z

@cursor review

cursor

Stale Bugbot comment from a previous run.

oscarkey

hey! One thing re .to(), otherwise just comments on claude on comments.

I'm wondering if another option might be to replace the current cache with an extension of tabpfn.inference._PerDeviceModelCache, which might allow us to neatly have .to() work and properly support multi-device inference (atm the model is only cached on one device?). But, the current approach seems pragmatic as it's not part of the public api.

oscarkey · 2026-09-02T08:48:14Z

+# Sharing is what a caller who enables it takes on: two estimators served one
+# entry hold the same module, so moving or training either reaches the other.
+#
+# `cache_trainset_representation` is not part of the key: every architecture's


Maybe: We don't include the cache_trainset_representation kwarg of ArchitectureModule.get_architecture() in the key, because it is no longer used since InferenceEngineCacheKV was removed.

I now went for the real fix and opened a PR to entirely wipe cache_trainset_representation #1250

I rebased the current PR on top of it.

oscarkey · 2026-09-02T11:48:07Z

+# holds a classifier and a regressor. Entries are shared by reference and left in
+# ``eval()`` mode, for repeated sequential fit/predict.
+#
+# Sharing is what a caller who enables it takes on: two estimators served one


What do you think about disabling .to() when the cache is enabled? Because the key is based on the initialisation options, not the actual device, you can get some confusing behaviour e.g.

a = TabPFNClassifier(device="cpu").to("cuda") b = TabPFNClassifier(device="cpu") # actually on cuda

oh no wait I'm wrong! Actually both a and b will end up on cpu. So that's mostly covered by this comment already.

so do you think we are fine as is?

jmkuebler · 2026-09-08T09:40:48Z

@cursor review

cursor

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit 6aebfda. Configure here.}

The cache (opt-in via `TABPFN_MODEL_CACHE_SIZE`) hands one built model to every estimator that asks for it. Estimators then place what they are given: the inference engine moves the module with an in-place `model.to(device)` and casts it with `model.type(dtype)`. Neither was part of the cache key, so two estimators wanting different placements were served the same instance: Placeholder storage has not been allocated on MPS device! mat1 and mat2 must have the same dtype, but got Float and Half Both are now in the key. The dtype cannot be repaired instead: casting to float16 is lossy, so a model that has been cast can never serve a full-precision caller. `has_equal_borders` compares values rather than requiring a shared device, since a cached criterion carries the device of the fit that placed it. `fit_mode="fit_with_cache"` is covered too. It was excluded on the grounds that `cache_trainset_representation` makes the module accumulate the train-set representation, but every architecture's `get_architecture` deletes that flag and the KV cache is owned by the inference engine. Setting `TABPFN_MODEL_CACHE_SIZE` over tests/test_classifier_interface.py and tests/test_regressor_interface.py fails 38 tests without this and passes with it, taking those two files from 146s to 117s; over the whole suite, 398s to 368s. The default stays off, and two estimators served one entry still share the module, so moving or training either reaches the other. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Co-authored-by: Oscar <oscar@priorlabs.ai>

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…d by #1250 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

jmkuebler changed the title ~~fix: key the built-model cache on placement, and enable it by default~~ key the built-model cache on placement, and enable it by default Aug 31, 2026

jmkuebler marked this pull request as ready for review August 31, 2026 15:00

cursor Bot reviewed Aug 31, 2026

View reviewed changes

Comment thread src/tabpfn/classifier.py Outdated

claude Bot reviewed Aug 31, 2026

View reviewed changes

Comment thread src/tabpfn/model_loading.py

claude Bot reviewed Aug 31, 2026

View reviewed changes

Comment thread src/tabpfn/classifier.py Outdated

cursor Bot reviewed Aug 31, 2026

View reviewed changes

Comment thread src/tabpfn/classifier.py Outdated

cursor Bot reviewed Sep 1, 2026

View reviewed changes

Comment thread src/tabpfn/base.py Outdated

cursor Bot reviewed Sep 1, 2026

View reviewed changes

Comment thread src/tabpfn/base.py Outdated

jmkuebler marked this pull request as draft September 1, 2026 08:32

jmkuebler changed the title ~~key the built-model cache on placement, and enable it by default~~ feat: cache built models for fit_mode="fit_with_cache" too Sep 1, 2026

jmkuebler changed the title ~~feat: cache built models for fit_mode="fit_with_cache" too~~ fix: key the built-model cache on placement, and cover fit_with_cache Sep 1, 2026

jmkuebler force-pushed the jonas/res_2743_model_cache_placement branch from 86adfcc to a977596 Compare September 1, 2026 09:16

jmkuebler marked this pull request as ready for review September 1, 2026 09:22

cursor Bot reviewed Sep 1, 2026

View reviewed changes

jmkuebler requested review from ggprior and oscarkey September 1, 2026 09:41

oscarkey reviewed Sep 2, 2026

View reviewed changes

jmkuebler mentioned this pull request Sep 8, 2026

Remove the dead cache_trainset_representation flag #1250

Open

7 tasks

jmkuebler changed the base branch from main to jonas/drop_cache_trainset_representation September 8, 2026 09:12

jmkuebler force-pushed the jonas/res_2743_model_cache_placement branch from 5bb9201 to 6aebfda Compare September 8, 2026 09:12

jmkuebler changed the title ~~fix: key the built-model cache on placement, and cover fit_with_cache~~ fix: key the built-model cache on device and inference dtype Sep 8, 2026

cursor Bot reviewed Sep 8, 2026

View reviewed changes

jmkuebler and others added 2 commits September 8, 2026 12:44

Apply batched suggestions from code review

74769eb

Co-authored-by: Oscar <oscar@priorlabs.ai>

jmkuebler and others added 2 commits September 8, 2026 12:44

Flatten the cache key and reword the flag note per review

f4467ce

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Drop the flag note and the fit_with_cache changelog entry, now covere…

e17adc0

…d by #1250 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

jmkuebler force-pushed the jonas/res_2743_model_cache_placement branch from 6aebfda to e17adc0 Compare September 8, 2026 12:44

Uh oh!

Conversation

jmkuebler commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Issue

Motivation and Context

The suite is the clearest evidence

What a caller still owns

Public API Changes

How Has This Been Tested?

Checklist

Uh oh!

jmkuebler commented Aug 31, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jmkuebler commented Sep 1, 2026

Uh oh!

Uh oh!

Uh oh!

jmkuebler commented Sep 1, 2026

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

oscarkey left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

oscarkey Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

jmkuebler Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

oscarkey Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

oscarkey Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

jmkuebler Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

jmkuebler commented Sep 8, 2026

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jmkuebler commented Aug 31, 2026 •

edited

Loading

cursor Bot left a comment •

edited

Loading