Build models from a checkpoint without the random parameter initialisation

[Innixma]

Stacked on #1255.

_build_model constructs the architecture and then loads the checkpoint with a strict load_state_dict, which overwrites every parameter and persistent buffer. The construction still ran every layer's reset_parameters, drawing random weights only to discard them: about a quarter of a second per build on a full-size checkpoint, paid at every fit and, under bagging, at every child.

The architecture is now built inside a small context manager that turns the torch.nn.init functions into no-ops for the duration and restores them afterwards, so the tensors are allocated but not filled before the load. The result is the same model: the architectures register no non-persistent buffers, and the tensors they initialise themselves are parameters, so all of them come from the checkpoint. A test loads small v2 and v3 checkpoints and checks that every named parameter and buffer equals the source model's exactly; a second test checks the patch is scoped to the block.

On an 83.5M-parameter checkpoint the build goes from 287 ms to 19 ms, most of the remainder being the state-dict copy; the reference-prediction consistency tests pass unchanged. One caveat for reviewers: the context manager patches module-level functions, so a model constructed concurrently in another thread of the same process during that window would also skip its initialisation. The library builds models on the calling thread, and the parallel paths use separate processes, so no such window exists in the code as it stands.

🤖 Generated with Claude Code

https://claude.ai/code/session_01ELdutHiUqkvynzEP7EnPsi

[oscarkey]

hey! Sorry for the drive-by comment, I just happened to be discussing this in the office yesterday and then spotted your PR. I'm just wondering if you also tried PyTorch's suggested approach of using the meta device? That won't affect other threads in the process and won't miss future initialisation functions, but maybe it doesn't work for us.

[Innixma]

Below is from Claude, but TLDR it was actually a bit slower (and adds 750ms cold start time) and also unsafe to do without adding guards around it. Do we have concerns about other threads in the process for this scenario? I'm unsure how we would test all the cases.

Claude:

Thanks, good pointer. I tried it. Same 52 small datasets (23 binary, 16 multiclass, 13 regression), a single estimator fit once per dataset through the AutoGluon wrapper on one GPU, every model build timed in-process:

build	fit total, 52 fits	per-build median	first build (cold, incl. checkpoint load)
skip-init (this PR)	15.2 s	20.9 ms	549 ms
meta + load_state_dict(assign=True)	crashed on the first regression dataset
meta + assign + guards	15.9 s	25.3 ms	1285 ms

Predictions are bit-identical between skip-init and the guarded meta build on all 52 datasets. The build is about 2 s of the 15 s either way, and the 0.7 s gap between the two runs is within run-to-run noise (a repeat of the skip-init run the day before came out at 14.6 s).

The bare meta variant fails with Tensor.item() cannot be called on meta tensors inside FullSupportBarDistribution.__init__: one of our architectures keeps a plain tensor attribute that aliases its regression-borders buffer. assign=True swaps the buffer for the checkpoint tensor, but the alias still points at the meta tensor, and _resolve_regression_borders hands that to the criterion. The "guards" variant walks the modules after loading, re-points any meta attribute at the loaded tensor of the same name and shape, and asserts nothing on meta remains. That works, but it is a name-matching heuristic and a second mechanism to maintain. assign=True also makes the parameters alias the cached state dict, so on the CPU path any in-place mutation of a model would corrupt the cache for later estimators.

In a microbenchmark the meta build was ~2x faster (about 15 ms vs 30 ms per build), but that did not survive the real path. I agree with both of your points in favour of meta (no process-wide patch, no list of init functions to keep current); given no end-to-end gain and the two new failure modes, I'll stick with skip-init here. Worth revisiting if host memory during builds becomes the constraint, since assign avoids the second copy of the weights.

[oscarkey]

ah thank you this is very useful! If we don't need to get this in urgently, I suggest we fix the architecture so it works with the meta device. From a very quick look, I think we need to make this line an @property. This will then mean we don't need the guard function, which I'd hope will shave off a few ms and bring us close to the patching approach. Although, I'm not sure what's going on with the cold-start time.

For me, I'd prefer to use the recommended approach even if it is a bit slower, as I'm worried that patching creates several tricky latent bugs:

• Anything that takes a reference to one of the random functions while the patch is applied will keep the patched one forever (although you could have the patches raise an error if used after the patch is removed)
• Anything that takes a reference before the patch is applied doesn't get the patch
• If anyone does use TabPFN with multi-threading, we'll potentially silently and randomly break their results

For people who want to go really fast, we can help them use the TABPFN_MODEL_CACHE_SIZE option.

What do you think?