Removing the PyTorch dependency - how should we handle existing model.pt files?

[jmwright]

I’ve been chipping away at removing the PyTorch dependency from Monty so it can run as a lighter-weight runtime on more constrained hardware. The recent pieces (de-torching the point-pair-feature math and swapping torch.multiprocessing for stdlib multiprocessing) were straightforward refactors. The next big piece is the object-model data structures (graphs → NumPy, sparse voxel grids → SciPy/NumPy), which I think also “changes shape without changing meaning.”

But there’s one part I don’t want to contribute without discussion, because it could introduce breaking behavior for users via model persistence. Right now Monty saves/loads models with torch.save / torch.load(…, weights_only=False) - which is really just pickle with special tensor handling. To drop the torch dependency entirely, we have to stop using it for model.pt too. The problem is that people (myself included) have pretrained models on disk, and I am guessing it is not acceptable to break them.

The direction I’m leaning toward:

1. Write new models with stdlib pickle (safe once the model objects are themselves torch-free - after some upcoming data-structure work).
2. Add a format-detecting loader: new pickle files load directly and legacy PyTorch-written model.pt files load through a small torch-free compatibility shim that reads torch’s archive format and rebuilds the tensors as NumPy arrays - no torch import needed.
3. Migrate TBP’s officially distributed pretrained models once so the documented path never needs the shim. The shim is purely for everyone’s older/own models.

That way torch can finally leave pyproject, and nobody’s existing models break.

So, a few questions:

• Is this RFC-worthy? Changing the model file format feels like a place to pause for discussion, so my instinct is yes. However, if the user never sees the change because of the compatibility layer, and if Monty’s end-result operation does not fundamentally change, maybe not.
• Should we move to a new save format: pickle, or something safer? Pickle is a drop-in for what torch.load(weights_only=False) already does (same arbitrary-code-execution caveat - no regression), but if we’re touching this anyway, is now the time to move to a safe, framework-agnostic format (safetensors-style arrays + a manifest)? That’s a bigger lift.
• Do we keep the .pt filename for continuity (the loader can sniff the format regardless), or rename to something that isn’t a PyTorch misnomer?
• How long should the shim live, and which torch save-format versions must it support - just the current reference version, or a range?

I’m not to the point where I’m ready to do this work quite yet. I wanted to start this discussion before I got there.

Please let me know what you think.

[tslominski]

Hi @jmwright, thank you for removing PyTorch where we don’t need it. You already addressed two contributing causes with your work so far .

@jmwright

and I am guessing it is not acceptable to break them.

Break them . Models will need to be retrained. Monty is not version 1, it is version 0, which can break anything from release to release.

It is worth noting that the MAJOR version number 0 is special. When MAJOR is 0, then incrementing the MINOR version number can include backwards-incompatible changes. MAJOR version number 0 is the usual way of indicating that the software is still in development and not ready for production use.

We will use MAJOR version number 0 for the ongoing development of Monty. Incrementing the MAJOR version number beyond 0 should be a future RFC, as it will be a significant event.

– tbp.monty/rfcs/0007_monty_versioning.md at main · thousandbrainsproject/tbp.monty · GitHub

I think we might still need a discussion/RFC on what the non-PyTorch format should be. But does this help remove some of the constraints?

Here’s what we’re currently tracking on the Current Reality Tree regarding model loading. Something worth highlighting, which might not be obvious, is that we embed tbp.monty.frameworks.models.object_model.GraphObjectModel|GridObjectModel, which requires tbp.monty to be available for loading.

[hlee]

Just throwing my two cents out there for an option for non-PyTorch format: safetensors

Despite the name “tensor” that reminds us of ML and support by PyTorch, it is the modern replacement for torch.save because using this can cause arbitrary code execution and this is the “safe” version of that (link). It does not need torch dependency.