Prevent _legacy_load with weights_only=True (#144993)

Prevent _legacy_load with weights_only=True (#144914)

Pull Request resolved: #144914
Approved by: https://github.com/malfet, https://github.com/albanD

(cherry picked from commit 7c3aa1d)

Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>

test/quantization/bc/test_backward_compatibility.py

@@ -112,12 +112,14 @@ def _test_op(
             torch.jit.save(torch.jit.trace(qmodule, input_tensor), traced_module_file)
             torch.save(qmodule(input_tensor), expected_file)
 
-        input_tensor = torch.load(input_file)
+        # weights_only=False as file was saved in .tar format
+        input_tensor = torch.load(input_file, weights_only=False)
         # weights_only = False as sometimes get ScriptObject here
         qmodule.load_state_dict(torch.load(state_dict_file, weights_only=False))
         qmodule_scripted = torch.jit.load(scripted_module_file)
         qmodule_traced = torch.jit.load(traced_module_file)
-        expected = torch.load(expected_file)
+        # weights_only=False as file was saved in .tar format
+        expected = torch.load(expected_file, weights_only=False)
         self.assertEqual(qmodule(input_tensor), expected, atol=prec)
         self.assertEqual(qmodule_scripted(input_tensor), expected, atol=prec)
         self.assertEqual(qmodule_traced(input_tensor), expected, atol=prec)

test/test_serialization.py

@@ -224,9 +224,6 @@ def _test_serialization(self, weights_only):
     def test_serialization(self):
         self._test_serialization(False)
 
-    def test_serialization_safe(self):
-        self._test_serialization(True)
-
     def test_serialization_filelike(self):
         # Test serialization (load and save) with a filelike object
         b = self._test_serialization_data()
@@ -362,9 +359,6 @@ def _test_serialization(conversion):
     def test_serialization_sparse(self):
         self._test_serialization(False)
 
-    def test_serialization_sparse_safe(self):
-        self._test_serialization(True)
-
     def test_serialization_sparse_invalid(self):
         x = torch.zeros(3, 3)
         x[1][1] = 1
@@ -510,9 +504,6 @@ def __reduce__(self):
     def test_serialization_backwards_compat(self):
         self._test_serialization_backwards_compat(False)
 
-    def test_serialization_backwards_compat_safe(self):
-        self._test_serialization_backwards_compat(True)
-
     def test_serialization_save_warnings(self):
         with warnings.catch_warnings(record=True) as warns:
             with tempfile.NamedTemporaryFile() as checkpoint:
@@ -557,7 +548,8 @@ def load_bytes():
         def check_map_locations(map_locations, dtype, intended_device):
             for fileobject_lambda in fileobject_lambdas:
                 for map_location in map_locations:
-                    tensor = torch.load(fileobject_lambda(), map_location=map_location)
+                    # weigts_only=False as the downloaded file path uses the old serialization format
+                    tensor = torch.load(fileobject_lambda(), map_location=map_location, weights_only=False)
 
                     self.assertEqual(tensor.device, intended_device)
                     self.assertEqual(tensor.dtype, dtype)
@@ -600,7 +592,8 @@ def test_load_nonexistent_device(self):
 
         error_msg = r'Attempting to deserialize object on a CUDA device'
         with self.assertRaisesRegex(RuntimeError, error_msg):
-            _ = torch.load(buf)
+            # weights_only=False as serialized is in legacy format
+            _ = torch.load(buf, weights_only=False)
 
     @unittest.skipIf((3, 8, 0) <= sys.version_info < (3, 8, 2), "See https://bugs.python.org/issue39681")
     def test_serialization_filelike_api_requirements(self):
@@ -720,7 +713,8 @@ def test_serialization_storage_slice(self):
                       b'\x00\x00\x00\x00')
 
         buf = io.BytesIO(serialized)
-        (s1, s2) = torch.load(buf)
+        # serialized was saved with PyTorch 0.3.1
+        (s1, s2) = torch.load(buf, weights_only=False)
         self.assertEqual(s1[0], 0)
         self.assertEqual(s2[0], 0)
         self.assertEqual(s1.data_ptr() + 4, s2.data_ptr())
@@ -837,6 +831,24 @@ def wrapper(*args, **kwargs):
     def __exit__(self, *args, **kwargs):
         torch.save = self.torch_save
 
+
+# used to set weights_only=False in _use_new_zipfile_serialization=False tests
+class load_method:
+    def __init__(self, weights_only):
+        self.weights_only = weights_only
+        self.torch_load = torch.load
+
+    def __enter__(self, *args, **kwargs):
+        def wrapper(*args, **kwargs):
+            kwargs['weights_only'] = self.weights_only
+            return self.torch_load(*args, **kwargs)
+
+        torch.load = wrapper
+
+    def __exit__(self, *args, **kwargs):
+        torch.load = self.torch_load
+
+
 Point = namedtuple('Point', ['x', 'y'])
 
 class ClassThatUsesBuildInstruction:
@@ -873,14 +885,25 @@ def test(f_new, f_old):
 
             torch.save(x, f_old, _use_new_zipfile_serialization=False)
             f_old.seek(0)
-            x_old_load = torch.load(f_old, weights_only=weights_only)
+            x_old_load = torch.load(f_old, weights_only=False)
             self.assertEqual(x_old_load, x_new_load)
 
         with AlwaysWarnTypedStorageRemoval(True), warnings.catch_warnings(record=True) as w:
             with tempfile.NamedTemporaryFile() as f_new, tempfile.NamedTemporaryFile() as f_old:
                 test(f_new, f_old)
             self.assertTrue(len(w) == 0, msg=f"Expected no warnings but got {[str(x) for x in w]}")
 
+    def test_old_serialization_fails_with_weights_only(self):
+        a = torch.randn(5, 5)
+        with BytesIOContext() as f:
+            torch.save(a, f, _use_new_zipfile_serialization=False)
+            f.seek(0)
+            with self.assertRaisesRegex(
+                RuntimeError,
+                "Cannot use ``weights_only=True`` with files saved in the .tar format used before version 1.6."
+            ):
+                torch.load(f, weights_only=True)
+
 
 class TestOldSerialization(TestCase, SerializationMixin):
     # unique_key is necessary because on Python 2.7, if a warning passed to
@@ -956,8 +979,7 @@ def test_serialization_offset(self):
         self.assertEqual(i, i_loaded)
         self.assertEqual(j, j_loaded)
 
-    @parametrize('weights_only', (True, False))
-    def test_serialization_offset_filelike(self, weights_only):
+    def test_serialization_offset_filelike(self):
         a = torch.randn(5, 5)
         b = torch.randn(1024, 1024, 512, dtype=torch.float32)
         i, j = 41, 43
@@ -969,16 +991,16 @@ def test_serialization_offset_filelike(self, weights_only):
             self.assertTrue(f.tell() > 2 * 1024 * 1024 * 1024)
             f.seek(0)
             i_loaded = pickle.load(f)
-            a_loaded = torch.load(f, weights_only=weights_only)
+            a_loaded = torch.load(f)
             j_loaded = pickle.load(f)
-            b_loaded = torch.load(f, weights_only=weights_only)
+            b_loaded = torch.load(f)
         self.assertTrue(torch.equal(a, a_loaded))
         self.assertTrue(torch.equal(b, b_loaded))
         self.assertEqual(i, i_loaded)
         self.assertEqual(j, j_loaded)
 
     def run(self, *args, **kwargs):
-        with serialization_method(use_zip=False):
+        with serialization_method(use_zip=False), load_method(weights_only=False):
             return super().run(*args, **kwargs)
 
 

torch/serialization.py

@@ -1482,15 +1482,10 @@ def _get_wo_message(message: str) -> str:
                 "please torch.save your checkpoint with this option in order to use mmap."
             )
         if weights_only:
-            try:
-                return _legacy_load(
-                    opened_file,
-                    map_location,
-                    _weights_only_unpickler,
-                    **pickle_load_args,
-                )
-            except pickle.UnpicklingError as e:
-                raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
+            raise RuntimeError(
+                "Cannot use ``weights_only=True`` with files saved in the "
+                ".tar format used before version 1.6. " + UNSAFE_MESSAGE
+            )
         return _legacy_load(
             opened_file, map_location, pickle_module, **pickle_load_args
         )