Revert "Let a serialized fake-tensor state_dict be read back (#22111)"

exir/serde/export_serialize.py

@@ -17,7 +17,6 @@
 import heapq
 import inspect
 import io
-import contextlib
 import json
 import logging
 import math
@@ -305,16 +304,7 @@ def _reconstruct_fake_tensor(
     assert len(_CURRENT_DESERIALIZER) != 0, "Need access to current deserializer state"
     fake_tensor = _CURRENT_DESERIALIZER[-1].deserialize_tensor_meta(tensor_meta)
     if is_parameter:
-        # `nn.Parameter` defaults `requires_grad` to True, which raises on the
-        # integer dtypes quantized weights use. Derived rather than recorded by
-        # the reducer, whose payload is the on-disk format: a float parameter
-        # that had it False therefore comes back True.
-        requires_grad = (
-            fake_tensor.dtype.is_floating_point or fake_tensor.dtype.is_complex
-        )
-        fake_tensor = torch.nn.Parameter(  # type: ignore[assignment]
-            fake_tensor, requires_grad=requires_grad
-        )
+        fake_tensor = torch.nn.Parameter(fake_tensor)  # type: ignore[assignment]
     return fake_tensor
 
 
@@ -346,19 +336,7 @@ def deserialize_torch_artifact(
         return {}
     buffer = io.BytesIO(serialized)
     buffer.seek(0)
-    # A fake tensor reduces to `_reconstruct_fake_tensor`, which
-    # `weights_only=True` refuses unless it is allowlisted.
-    #
-    # Conditional because `safe_globals` subtracts from one process-global set on
-    # exit, with no refcount: entering it when the caller has already allowed
-    # this global would revoke their registration.
-    allowance = (
-        contextlib.nullcontext()
-        if _reconstruct_fake_tensor in torch.serialization.get_safe_globals()
-        else torch.serialization.safe_globals([_reconstruct_fake_tensor])
-    )
-    with allowance:
-        artifact = torch.load(buffer, weights_only=True)
+    artifact = torch.load(buffer, weights_only=True)
     assert isinstance(artifact, (tuple, dict))
     return artifact
 

exir/tests/test_serde.py

@@ -9,8 +9,7 @@
 import io
 import tempfile
 import unittest
-from contextlib import contextmanager
-from typing import Dict, Iterator, Tuple
+from typing import Tuple
 
 import executorch.exir as exir
 
@@ -29,19 +28,10 @@
     EdgeProgramManager,
     to_edge_transform_and_lower,
 )
-from executorch.exir.serde.export_serialize import (
-    _CURRENT_DESERIALIZER,
-    _reconstruct_fake_tensor,
-    deserialize_torch_artifact,
-    GraphModuleDeserializer,
-    serialize_torch_artifact,
-)
 from executorch.exir.serde.serialize import deserialize, serialize
 from torch import nn
-from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
 from torch.export import export
 from torch.export.exported_program import ExportedProgram as TorchExportedProgram
-from torch.export.graph_signature import InputKind
 from torch.utils import _pytree as pytree
 
 
@@ -360,154 +350,3 @@ def forward(self, x, y):
             torch.randn(10, 10),
         )
         self.check_serde(MemoryOpsModule(), inputs)
-
-
-class TestFakeTensorArtifacts(unittest.TestCase):
-    """`state_dict` / `constants` entries may be FakeTensors.
-
-    `_reduce_fake_tensor` writes one as its `TensorMeta` -- shape and dtype, no
-    storage -- so a caller can record a program's structure without its weight
-    values. The tests above only ever serialize real tensors.
-    """
-
-    @contextmanager
-    def _deserializer_on_the_stack(self) -> Iterator[None]:
-        """Push what `_reconstruct_fake_tensor` reads.
-
-        It rebuilds tensors through the current deserializer's fake mode, which
-        only `deserialize()` puts on this stack. These tests round-trip the
-        artifact alone, so they push it themselves.
-        """
-        deserializer = GraphModuleDeserializer()
-        deserializer.fake_tensor_mode = FakeTensorMode()
-        _CURRENT_DESERIALIZER.append(deserializer)
-        try:
-            yield
-        finally:
-            _CURRENT_DESERIALIZER.pop()
-
-    def _round_trip(self, artifact: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        with self._deserializer_on_the_stack():
-            return deserialize_torch_artifact(serialize_torch_artifact(artifact))
-
-    def test_reading_leaves_a_callers_own_allowlist_entry_alone(self) -> None:
-        """Reading must not revoke an allowance the caller set up itself.
-
-        `safe_globals` subtracts on exit with no refcount, so entering it
-        unconditionally removes a registration made with `add_safe_globals`.
-        """
-        previous = torch.serialization.get_safe_globals()
-        torch.serialization.add_safe_globals([_reconstruct_fake_tensor])
-        try:
-            self._round_trip({"w": FakeTensorMode().from_tensor(torch.randn(2))})
-            self.assertIn(
-                _reconstruct_fake_tensor,
-                torch.serialization.get_safe_globals(),
-                "the caller's own registration must survive the load",
-            )
-        finally:
-            torch.serialization.clear_safe_globals()
-            torch.serialization.add_safe_globals(previous)
-
-    def test_fake_tensor_artifact_round_trips(self) -> None:
-        """The read path must accept what the write path produces.
-
-        `weights_only=True` refuses any global that is not allowlisted,
-        `_reconstruct_fake_tensor` included.
-        """
-        fake_mode = FakeTensorMode()
-        artifact = {"w": fake_mode.from_tensor(torch.randn(4, 8))}
-        self.assertIsInstance(artifact["w"], FakeTensor, "precondition")
-
-        restored = self._round_trip(artifact)
-
-        self.assertEqual({"w"}, set(restored))
-        self.assertIsInstance(restored["w"], FakeTensor)
-        self.assertEqual(torch.Size([4, 8]), restored["w"].shape)
-        self.assertEqual(torch.float32, restored["w"].dtype)
-
-    def test_integer_parameter_survives_deserialization(self) -> None:
-        """A parameter of integer dtype must round-trip.
-
-        Quantized weights are ones. Rebuilding with `nn.Parameter`'s default of
-        `requires_grad=True` raises `only Tensors of floating point dtype can
-        require gradients`.
-        """
-        fake_mode = FakeTensorMode()
-        codes = torch.nn.Parameter(
-            torch.randint(0, 255, (4, 8), dtype=torch.uint8), requires_grad=False
-        )
-        artifact = {"codes": fake_mode.from_tensor(codes)}
-        self.assertIsInstance(artifact["codes"], torch.nn.Parameter, "precondition")
-
-        restored = self._round_trip(artifact)
-
-        self.assertIsInstance(restored["codes"], FakeTensor)
-        self.assertIsInstance(restored["codes"], torch.nn.Parameter)
-        self.assertEqual(torch.uint8, restored["codes"].dtype)
-        self.assertFalse(restored["codes"].requires_grad)
-
-    def test_requires_grad_is_derived_from_the_dtype(self) -> None:
-        """`requires_grad` is rebuilt, not recovered -- pinned so it stays known.
-
-        The reducer records only *that* an entry was a parameter, so a float one
-        that had `requires_grad=False` comes back True.
-        """
-        fake_mode = FakeTensorMode()
-        artifact = {
-            "float_param": fake_mode.from_tensor(
-                torch.nn.Parameter(torch.randn(4), requires_grad=False)
-            ),
-            "int_param": fake_mode.from_tensor(
-                torch.nn.Parameter(
-                    torch.zeros(4, dtype=torch.int64), requires_grad=False
-                )
-            ),
-        }
-
-        restored = self._round_trip(artifact)
-
-        self.assertTrue(restored["float_param"].requires_grad, "derived, not kept")
-        self.assertFalse(restored["int_param"].requires_grad, "cannot be anything else")
-
-    def test_program_with_a_fake_state_dict_round_trips(self) -> None:
-        """The whole program, not just the artifact -- the shape callers use.
-
-        Only this reaches `_verify_exported_program_signature`, which requires
-        every `InputKind.PARAMETER` entry to still be an `nn.Parameter`.
-        """
-
-        class Quantized(nn.Module):
-            def __init__(self) -> None:
-                super().__init__()
-                self.codes = nn.Parameter(
-                    torch.randint(0, 255, (4, 4), dtype=torch.uint8),
-                    requires_grad=False,
-                )
-                self.scales = nn.Parameter(torch.rand(4, 4), requires_grad=False)
-
-            def forward(self, x: torch.Tensor) -> torch.Tensor:
-                return x + self.codes.to(torch.float32) * self.scales
-
-        program = to_edge(
-            export(Quantized(), (torch.randn(4, 4),), strict=True)
-        ).exported_program()
-        self.assertIn(
-            InputKind.PARAMETER,
-            [spec.kind for spec in program.graph_signature.input_specs],
-            "the model must have parameters for this to test anything",
-        )
-
-        fake_mode = FakeTensorMode()
-        program._state_dict = {
-            name: fake_mode.from_tensor(tensor, static_shapes=True)
-            for name, tensor in program.state_dict.items()
-        }
-
-        restored = deserialize(serialize(program))
-
-        self.assertEqual(set(program.state_dict), set(restored.state_dict))
-        for name, tensor in restored.state_dict.items():
-            self.assertIsInstance(tensor, FakeTensor, f"{name} should stay fake")
-            self.assertIsInstance(tensor, nn.Parameter, f"{name} stays a Parameter")
-        self.assertEqual(torch.uint8, restored.state_dict["codes"].dtype)