Repository: pytorch
Issue: Export-based GraphModules fail with deepcopy
Labels: oncall: export

### 🐛 Describe the bug


Some export-based `GraphModule`s fail with deepcopy, for example, ResNet, MobileNet, MNASNet, ShuffleNet. The error message is: `NotImplementedError: argument of type: <class 'torch.fx.graph_module.GraphModule.__new__.<locals>.GraphModuleImpl'>`. This operation was successful in PyTorch 2.1 but failed in PyTorch 2.2.

```python
import torchvision
import torch
from copy import deepcopy

model = torchvision.models.resnet18(weights=None)
input_data = torch.rand(1, 3, 224, 224)
gm = torch._export.capture_pre_autograd_graph(model, (input_data,))
deepcopy(gm)
```

full trace:

```python
[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in deepcopy(x, memo, _nil)
    151             copier = getattr(x, "__deepcopy__", None)
    152             if copier is not None:
--> 153                 y = copier(memo)
    154             else:
    155                 reductor = dispatch_table.get(cls)

[/usr/local/lib/python3.10/dist-packages/torch/fx/graph_module.py](https://localhost:8080/#) in __deepcopy__(self, memo)
    791         res = type(self).__new__(type(self))
    792         memo[id(self)] = res
--> 793         fake_mod = _CodeOnlyModule(copy.deepcopy(self.__dict__, memo))
    794         self._deepcopy_init()(res, fake_mod, fake_mod.__dict__["_graph"])
    795         # hooks are lost during `GraphModule.__init__`, so we need to copy over

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in deepcopy(x, memo, _nil)
    144     copier = _deepcopy_dispatch.get(cls)
    145     if copier is not None:
--> 146         y = copier(x, memo)
    147     else:
    148         if issubclass(cls, type):

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in _deepcopy_dict(x, memo, deepcopy)
    229     memo[id(x)] = y
    230     for key, value in x.items():
--> 231         y[deepcopy(key, memo)] = deepcopy(value, memo)
    232     return y
    233 d[dict] = _deepcopy_dict

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in deepcopy(x, memo, _nil)
    170                     y = x
    171                 else:
--> 172                     y = _reconstruct(x, memo, *rv)
    173 
    174     # If is its own copy, don't memoize.

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    295             for key, value in dictiter:
    296                 key = deepcopy(key, memo)
--> 297                 value = deepcopy(value, memo)
    298                 y[key] = value
    299         else:

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in deepcopy(x, memo, _nil)
    151             copier = getattr(x, "__deepcopy__", None)
    152             if copier is not None:
--> 153                 y = copier(memo)
    154             else:
    155                 reductor = dispatch_table.get(cls)

[/usr/local/lib/python3.10/dist-packages/torch/_tensor.py](https://localhost:8080/#) in __deepcopy__(self, memo)
    205                         setattr(new_tensor, slot, deepcopy(getattr(self, slot), memo))
    206 
--> 207             new_tensor.__dict__ = deepcopy(self.__dict__, memo)
    208 
    209             memo[id(self)] = new_tensor

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in deepcopy(x, memo, _nil)
    144     copier = _deepcopy_dispatch.get(cls)
    145     if copier is not None:
--> 146         y = copier(x, memo)
    147     else:
    148         if issubclass(cls, type):

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in _deepcopy_dict(x, memo, deepcopy)
    229     memo[id(x)] = y
    230     for key, value in x.items():
--> 231         y[deepcopy(key, memo)] = deepcopy(value, memo)
    232     return y
    233 d[dict] = _deepcopy_dict

[/usr/lib/python3.10/copy.py](https://localhost:8080/#) in deepcopy(x, memo, _nil)
    151             copier = getattr(x, "__deepcopy__", None)
    152             if copier is not None:
--> 153                 y = copier(memo)
    154             else:
    155                 reductor = dispatch_table.get(cls)

[/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py](https://localhost:8080/#) in __call__(self, *args, **kwargs)
    499 
    500     def __call__(self, *args, **kwargs):
--> 501         return self.tracer.create_proxy('call_method', self.attr, (self.root,) + args, kwargs)
    502 
    503 

[/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py](https://localhost:8080/#) in create_proxy(self, kind, target, args, kwargs, name, type_expr, proxy_factory_fn)
    183         '''
    184 
--> 185         args_ = self.create_arg(args)
    186         kwargs_ = self.create_arg(kwargs)
    187         assert isinstance(args_, tuple)

[/usr/local/lib/python3.10/dist-packages/torch/fx/experimental/proxy_tensor.py](https://localhost:8080/#) in create_arg(self, a)
    464             assert a.node.constant is not None
    465             return a.node.constant
--> 466         return super().create_arg(a)
    467 
    468     def unwrap_proxy(self, e):

[/usr/local/lib/python3.10/dist-packages/torch/fx/_symbolic_trace.py](https://localhost:8080/#) in create_arg(self, a)
    387             return self.create_node("get_attr", qualname, (), {})
    388 
--> 389         return super().create_arg(a)
    390 
    391     @compatibility(is_backward_compatible=True)

[/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py](https://localhost:8080/#) in create_arg(self, a)
    254             return type(a)(*args)  # type: ignore[arg-type]
    255         elif isinstance(a, (tuple, list)):
--> 256             return type(a)(self.create_arg(elem) for elem in a)
    257         elif isinstance(a, dict):
    258             r = {}

[/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py](https://localhost:8080/#) in <genexpr>(.0)
    254             return type(a)(*args)  # type: ignore[arg-type]
    255         elif isinstance(a, (tuple, list)):
--> 256             return type(a)(self.create_arg(elem) for elem in a)
    257         elif isinstance(a, dict):
    258             r = {}

[/usr/local/lib/python3.10/dist-packages/torch/fx/experimental/proxy_tensor.py](https://localhost:8080/#) in create_arg(self, a)
    464             assert a.node.constant is not None
    465             return a.node.constant
--> 466         return super().create_arg(a)
    467 
    468     def unwrap_proxy(self, e):

[/usr/local/lib/python3.10/dist-packages/torch/fx/_symbolic_trace.py](https://localhost:8080/#) in create_arg(self, a)
    387             return self.create_node("get_attr", qualname, (), {})
    388 
--> 389         return super().create_arg(a)
    390 
    391     @compatibility(is_backward_compatible=True)

[/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py](https://localhost:8080/#) in create_arg(self, a)
    269                 map_aggregate(k, no_node)
    270 
--> 271                 r[k] = self.create_arg(v)
    272             return r
    273         elif isinstance(a, slice):

[/usr/local/lib/python3.10/dist-packages/torch/fx/experimental/proxy_tensor.py](https://localhost:8080/#) in create_arg(self, a)
    464             assert a.node.constant is not None
    465             return a.node.constant
--> 466         return super().create_arg(a)
    467 
    468     def unwrap_proxy(self, e):

[/usr/local/lib/python3.10/dist-packages/torch/fx/_symbolic_trace.py](https://localhost:8080/#) in create_arg(self, a)
    387             return self.create_node("get_attr", qualname, (), {})
    388 
--> 389         return super().create_arg(a)
    390 
    391     @compatibility(is_backward_compatible=True)

[/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py](https://localhost:8080/#) in create_arg(self, a)
    290         elif isinstance(a, (*base_types, enum.Enum)) or a is None or a is ...:
    291             return a
--> 292         raise NotImplementedError(f"argument of type: {type(a)}")
    293 
    294     @compatibility(is_backward_compatible=True)

NotImplementedError: argument of type: <class 'torch.fx.graph_module.GraphModule.__new__.<locals>.GraphModuleImpl'>
```

### Versions



## Version
```bash
PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.9
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.1.58+-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             2
On-line CPU(s) list:                0,1
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) CPU @ 2.20GHz
CPU family:                         6
Model:                              79
Thread(s) per core:                 2
Core(s) per socket:                 1
Socket(s):                          1
Stepping:                           0
BogoMIPS:                           4399.99
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          32 KiB (1 instance)
L1i cache:                          32 KiB (1 instance)
L2 cache:                           256 KiB (1 instance)
L3 cache:                           55 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0,1
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Mitigation; PTE Inversion
Vulnerability Mds:                  Vulnerable; SMT Host state unknown
Vulnerability Meltdown:             Vulnerable
Vulnerability Mmio stale data:      Vulnerable
Vulnerability Retbleed:             Vulnerable
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Vulnerable
Vulnerability Spectre v1:           Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:           Vulnerable, IBPB: disabled, STIBP: disabled, PBRSB-eIBRS: Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Vulnerable

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] torch==2.2.1
[pip3] torchaudio==2.1.0+cu121
[pip3] torchdata==0.7.0
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.16.0
[pip3] torchvision==0.17.1
[pip3] triton==2.2.0
[conda] Could not collect
```

cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4