Propagate real tensor tracing with torchbind + fixing side effects (pytorch#138797)

Summary:
* Fixed real tensor tracing w/ torchbind objs by passing the cloned tensor obj. For now I just catch the exception and have an error message if the `_clone` fails, but up for discussion on what to do here
  * Separate question, should we require people to set up FakeScriptObjects and stuff for draft mode?
* Prevent side effects from happening when we do the first pass of custom ops profiling by cloning/copying everything. Not sure if deepcopying the model will succeed in all cases... But also I guess this path can be removed once custom ops profiling turns into one pass.

Test Plan: `buck2 run @//mode/dev-nosan //scripts/angelayi/draft_export:test_draft_export`

Reviewed By: ydwu4

Differential Revision: D64124825

Pull Request resolved: pytorch#138797
Approved by: https://github.com/ydwu4

Author: angelayi

torch/_library/fake_class_registry.py

@@ -1,9 +1,11 @@
 # mypy: allow-untyped-defs
+import copy
 import logging
 from typing import Any, Dict, Optional, Protocol, Tuple, Union
 
 import torch
 from torch._library.utils import parse_namespace
+from torch.utils._python_dispatch import _disable_current_modes
 
 
 log = logging.getLogger(__name__)
@@ -15,7 +17,18 @@ def __init__(self, wrapped_obj: Any, script_class_name: str, x: torch.ScriptObje
 
         # The fully qualified name of the class of original script object
         self.script_class_name = script_class_name
-        self.real_obj = x
+        try:
+            with _disable_current_modes():
+                self.real_obj = copy.deepcopy(x)
+        except RuntimeError:
+            log.warning(
+                "Unable to deepcopy the custom object %s. "
+                "Defaulting to the user given object. This might be "
+                "dangerous as side effects may be directly applied "
+                "to the object.",
+                script_class_name,
+            )
+            self.real_obj = x
 
 
 class FakeScriptMethod:

torch/_subclasses/fake_tensor.py

@@ -38,6 +38,7 @@
 import torch
 from torch import SymBool, SymFloat, SymInt, Tensor
 from torch._C._functorch import is_functorch_wrapped_tensor, is_legacy_batchedtensor
+from torch._library.fake_class_registry import FakeScriptObject
 from torch._prims_common import suggest_memory_format
 from torch._subclasses.meta_utils import (
     assert_eq,
@@ -1947,7 +1948,9 @@ def _dispatch_impl(
         args, kwargs = pytree.tree_unflatten(flat_args, args_spec)
         self.invalidate_written_to_constants(func, flat_arg_fake_tensors, args, kwargs)
 
-        def maybe_to_real_tensor(t: T) -> Optional[Union[T, Tensor]]:
+        def maybe_to_real_tensor(
+            t: T,
+        ) -> Optional[Union[T, Tensor, torch._C.ScriptObject]]:
             if isinstance(t, FakeTensor):
                 return t.real_tensor
             elif isinstance(t, py_sym_types):
@@ -1957,6 +1960,8 @@ def maybe_to_real_tensor(t: T) -> Optional[Union[T, Tensor]]:
                         self.shape_env.unbacked_var_to_val
                     )
                 )
+            elif isinstance(t, FakeScriptObject):
+                return t.real_obj
             else:
                 return t