BC-Breaking Change: torch.load is being flipped to use weights_only=True by default in the nightlies after #137602

frontend API
1

TL;DR

After warning of this change since version 2.4, #137602 has been merged which will change the default for the weights_only argument of torch.load from False to True in the nightlies (and version >= 2.6). See this note on torch.load with weights_only=True for a summary what weights_only=True does and all utilities related to weights_only.

Background

By default, torch.load uses weights_only=False which uses the pickle module provided by Python. This is inherently unsafe as the pickle module can import and execute arbitrary functions based on the bytecode within the pickle file.

A recent event that demonstrates this is here where a ByteDance intern attacked large cluster training via inserting remote code execution exploits into model checkpoints (among other things).

In contrast, the weights_only unpickler rewrites the stack machine of the unpickler to ensure that no modules/functions are imported and only certain allowlisted classes/functions (necessary for loading state_dicts of plain tensors) are called during the unpickling process.

This change is expected to incrementally improve the safety of torch.load.

What to expect

  • state_dicts of plain torch.Tensors should be loadable seamlessly
  • Loading serialized nn.Modules is against recommended best practices see here and will now explicitly fail, please set weights_only=False for these cases
  • Tensor subclasses, serialized numpy arrays and other objects will require the user to use safe_globals to iteratively allowlist certain classes/functions in the checkpoint (the error message when loading should give the appropriate instructions).

For more detailed instructions on all utilities related to weights_only please see this note on torch.load with weights_only=True.

You can also use TorchFix to find callsites where your code does not explicitly set weights_only.

What to do if you have issues

Do file an issue to pytorch/pytorch and tag me @mikaylagawarecki if you have issues.

5 Likes