`torch.load` planned default flip for `weights_only`

[mikaylagawarecki]

TL;DR

PyTorch is planning a BC-breaking change in torch.load to flip the default for weights_only from None (i.e. False) to True (and have added a warning to this effect in torch 2.4 :) ) that will break loading of tensor serialized when on XLA.

Context

Instead of using the default Unpickler provided by pickle, torch.load(weights_only=True) uses a custom Unpickler that restricts the allowed GLOBALs in the checkpoint (classes and functions) to those here (that required to build state_dicts).

The purpose of this is towards addressing the issue of remote code execution when using torch.load.

Another feature of this is that users can allowlist certain globals using add_safe_globals (in torch 2.4) or the safe_globals context manager (in torch nightly), a simple example being

import torch
from torch.serialization import safe_globals

class MyTensor(...):
     pass
 
 t = MyTensor(torch.randn(2, 3))    
 torch.save(t, "ckpt.pt")
 
 # This fails saying that MyTensor is not an allowed GLOBAL
 # t1 = torch.load("ckpt.pt", weights_only=True)
 
 # This succeeds
 with safe_globals([MyTensor]):
     torch.load("ckpt.pt", weights_only=True)

How this affects XLA

Notably, XLA uses a special path that uses numpy for serialization/deserialization see here. However, we have made a decision not to include the numpy GLOBALS required for unpickling in the defaut list as we do not control the codepaths numpy implements for pickling (see relevant GLOBALs here)

Ask

Opening this issue to figure out the best way to move forward re above to make the flip as smooth as possible!

Ideally, it would be good if the path for serializing XLA tensors could be refactored to not use numpy and we would definitely accept a PR that implements this!

Separately, for existing checkpoints I imagine there will be something that needs to be done there.

cc @JackCaoG

[JackCaoG]

@will-cromar any thoughts?

[mikaylagawarecki]

Gentle bump, @will-cromar what would be the best path forward here?

[JackCaoG]

Sorry @mikaylagawarecki for the delay, I will take a look next Monday.

[will-cromar]

Discussed this with @JackCaoG. We actually recommend against saving xla device tensors directly. We implement a wrapper that actually moves all of the data to CPU first:

xla/torch_xla/core/xla_model.py

Lines 1272 to 1274 in e657c87

	cpu_data = _maybe_convert_to_cpu(data, convert=should_write_data)
	if should_write_data:
	torch.save(cpu_data, file_or_path)

Is it possible to do the same thing in torch? That is, if you see an xla tensor, move it to the CPU and then it always loads as a CPU tensor.

I saw there are some other backends that use the same code path. Is there a plan to handle those? If you are planning add some interface we can implement for custom serialization/deserialization logic, it should also be possible for us to save/load data in an XLA-friendly format (via xla::Literal). I want to avoid upstreaming/inlining a bunch of XLA-specific logic, though.

[mikaylagawarecki]

Hey @will-cromar, thanks for the response here! First of all, I am very glad to hear that XLA recommends saving CPU tensors 😄 That means that the weights_only flip will be very minimally BC-breaking for XLA.

Is it possible to do the same thing in torch? That is, if you see an xla tensor, move it to the CPU and then it always loads as a CPU tensor.

It is definitely possible, but ~technically is also a BC-break within torch (separate from the weights_only BC break)

I chatted with @JackCaoG a bit offline, sharing the summary of what we came to here

If users always do tensor.cpu() before passing to torch.save, that is great and yea seems like the blast radius of the BC-break will be small for XLA! I will nevertheless be trying to remove the numpy dependency for 2.6 🙂
So after the torch.load weights_only flip from False->True`
(1) checkpoints with tensors on XLA saved in <2.6 will have errors like "numpy._reconstruct" is not a safe GLOBAL, please add it to safe_globals if you trust it
(2) checkpoints with tensors on XLA saved in >=2.6 will work seamlessly with torch.load(weights_only=True) (I will have a fix to remove numpy dependency)
(3) users who always move their XLA tensors to CPU before saving will not be affected at all

For the (perhaps tiny) portion of users who save their tensors without moving to CPU, I am personally inclined to keep the "if save on XLA, load on XLA" behavior (with the numpy dependency removed)

But if you feel strongly that always forcing CPU within torch.save is the better path that sounds plausible too