[Dynamo] Guard serialization for NN_MODULE (pytorch#152725)

Throws an error when attempting to serialize an NN_MODULE guard. It is not supported because it uses the unsupported ID_MATCH guard (pytorch#152330):

https://github.com/pytorch/pytorch/blob/a6dd1c2208f29a3169c1fe96bf4e79a10aa5647d/torch/_dynamo/guards.py#L1738-L1739

Pull Request resolved: pytorch#152725
Approved by: https://github.com/jansel

Author: jbschlosser

test/dynamo/test_guard_serialization.py

@@ -5,6 +5,7 @@
 import pickle
 import sys
 import types
+from unittest.mock import patch
 
 import torch
 import torch._dynamo.testing
@@ -604,6 +605,22 @@ def fn(x, y):
         # guard should fail for different y value
         self._test_check_fn(ref, loaded, {"x": torch.randn(3), "y": 6}, False)
 
+    def test_nn_module(self):
+        def fn(m, x):
+            return m(x)
+
+        m = GlobalModule()
+        x = torch.randn(3)
+
+        # config setting controls whether the NN_MODULE guard is installed
+        with patch("torch._dynamo.config.inline_inbuilt_nn_modules", False):
+            # we don't support NN_MODULE because it adds an ID_MATCH guard, and we don't
+            # support that in serialization
+            with self.assertRaisesRegex(
+                RuntimeError, "NN_MODULE guard cannot be serialized."
+            ):
+                self._test_serialization("NN_MODULE", fn, m, x)
+
     def test_dict_version(self):
         def fn(x):
             return pytree.tree_leaves(x)[0] + 1

torch/_dynamo/guards.py

@@ -1740,6 +1740,9 @@ def CONSTANT_MATCH(self, guard: Guard):
             self.EQUALS_MATCH(guard)
 
     def NN_MODULE(self, guard: Guard):
+        # don't support this in serialization because it uses unsupported ID_MATCH
+        if self.serialization_mode == "save":
+            raise RuntimeError("NN_MODULE guard cannot be serialized.")
         self.ID_MATCH(guard)
         val = self.get(guard.name)
         if hasattr(val, "training"):