Repository: pytorch
Issue: Dynamo recompiles code on every call when arguments include nn.Module
Labels: triaged,oncall: pt2,module: dynamo,release notes: dynamo

### 🐛 Describe the bug

When I `torch.compile` some code that receives an argument of type nn.Module, the recompilation is triggered on each call with a different instance. Expected it to recompile only some attributes relevant for the code are different.

Test code:

```python
import torch
import torch.nn as nn
from torch import Tensor

import logging
# torch._logging.set_logs(guards=True)
torch._logging.set_logs(recompiles=True, recompiles_verbose = True, fusion = True)

Mod = nn.Module

class A(Mod):
    def __init__(self, v:float):
        super().__init__()
        self.b = v
        
    def forward(self, x):
        return (x.sin()+self.b).sin()
       
def forward(net, x:Tensor) -> Tensor:
    return net.forward(x).sum()


dev = torch.device('cuda:0')

compile_args = dict(fullgraph=True, dynamic = True, backend = "inductor", options={'force_same_precision':True, 'disable_cpp_codegen':False, 'trace.graph_diagram':True, "triton.cudagraphs": False})
cforward = torch.compile(forward, **compile_args)

def compute(net):
    y = cforward(net, x)
    y.backward()

start_e = torch.cuda.Event(enable_timing=True)
end_e = torch.cuda.Event(enable_timing=True)
x = torch.rand((100,100,100), device =dev, requires_grad=True)
# for d in[1,2]:
#     torch._dynamo.mark_dynamic(x, d)

nets = [A(float(1.0)) for j in range(3)]
# nets = [nn.Sequential(*[A(float(j)) for j in range(3)]) for i in range(3)]

# WARMUP/ COMPILE
s = torch.cuda.Stream(device=dev)
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.device(dev):
    with torch.cuda.stream(s):  # warmup on a side stream, according to examples
        # check correctness
        for i in range(3):
            y = forward(nets[i], x)
            y.backward()
        # warmup compiled
        for i in range(3):
            print(f"_______________Net {i}_______________")    
            x.grad = None
            start_e.record()
            compute(nets[i])
            end_e.record()
            torch.cuda.synchronize()
            t = start_e.elapsed_time(end_e)
            print(f'Time: {t} ms')
            print(x.grad.mean())
            
    torch.cuda.current_stream().wait_stream(s)
    G = torch.cuda.CUDAGraph()
    with torch.cuda.graph(G):
        x.grad = None
        for i in range(3):
            compute(nets[i])

    for i in range(3):
        print(f"_______________Compiled Graphed Iteration {i}_______________")    
        tt = 0
        for j in range(100):
            start_e.record()
            G.replay()
            end_e.record()
            torch.cuda.synchronize()
            t = start_e.elapsed_time(end_e)
            tt += t
        print(f'Time: {tt/100} ms')
        print(x.grad.mean())
```

It prints
```
_______________Net 0_______________
V0902 15:57:19.824000 140183681587008 torch/_inductor/scheduler.py:1820] [0/0] [__fusion] ===== attempting fusion (1/10): 2 nodes =====
V0902 15:57:19.825000 140183681587008 torch/_inductor/scheduler.py:627] [0/0] [__fusion] cannot fuse buf0 with buf1: numel/rnumel mismatch (reduce) (123, 1), (((s0**3 + 122)//123), 123)
V0902 15:57:19.826000 140183681587008 torch/_inductor/scheduler.py:2133] [0/0] [__fusion] found 0 possible fusions
V0902 15:57:19.827000 140183681587008 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] completed fusion round (1/10): fused 2 nodes into 2 nodes
V0902 15:57:19.827000 140183681587008 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] 
V0902 15:57:19.827000 140183681587008 torch/_inductor/scheduler.py:1834] [0/0] [__fusion] ===== fusion complete (1 iterations) =====
V0902 15:57:21.398000 140183681587008 torch/_inductor/scheduler.py:1820] [0/0] [__fusion] ===== attempting fusion (1/10): 1 nodes =====
V0902 15:57:21.399000 140183681587008 torch/_inductor/scheduler.py:2133] [0/0] [__fusion] found 0 possible fusions
V0902 15:57:21.399000 140183681587008 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] completed fusion round (1/10): fused 1 nodes into 1 nodes
V0902 15:57:21.399000 140183681587008 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] 
V0902 15:57:21.399000 140183681587008 torch/_inductor/scheduler.py:1834] [0/0] [__fusion] ===== fusion complete (1 iterations) =====
Time: 4855.01953125 ms
tensor(0.1223, device='cuda:0')
_______________Net 1_______________
V0902 15:57:22.948000 140183681587008 torch/_inductor/scheduler.py:1820] [0/1] [__fusion] ===== attempting fusion (1/10): 2 nodes =====
V0902 15:57:22.949000 140183681587008 torch/_inductor/scheduler.py:627] [0/1] [__fusion] cannot fuse buf0 with buf1: numel/rnumel mismatch (reduce) (123, 1), (((s0**3 + 122)//123), 123)
V0902 15:57:22.950000 140183681587008 torch/_inductor/scheduler.py:2133] [0/1] [__fusion] found 0 possible fusions
V0902 15:57:22.950000 140183681587008 torch/_inductor/scheduler.py:1827] [0/1] [__fusion] completed fusion round (1/10): fused 2 nodes into 2 nodes
V0902 15:57:22.950000 140183681587008 torch/_inductor/scheduler.py:1827] [0/1] [__fusion] 
V0902 15:57:22.950000 140183681587008 torch/_inductor/scheduler.py:1834] [0/1] [__fusion] ===== fusion complete (1 iterations) =====
V0902 15:57:23.437000 140183681587008 torch/_inductor/scheduler.py:1820] [0/1] [__fusion] ===== attempting fusion (1/10): 1 nodes =====
V0902 15:57:23.437000 140183681587008 torch/_inductor/scheduler.py:2133] [0/1] [__fusion] found 0 possible fusions
V0902 15:57:23.437000 140183681587008 torch/_inductor/scheduler.py:1827] [0/1] [__fusion] completed fusion round (1/10): fused 1 nodes into 1 nodes
V0902 15:57:23.437000 140183681587008 torch/_inductor/scheduler.py:1827] [0/1] [__fusion] 
V0902 15:57:23.438000 140183681587008 torch/_inductor/scheduler.py:1834] [0/1] [__fusion] ===== fusion complete (1 iterations) =====
Time: 1210.2095947265625 ms
tensor(0.1223, device='cuda:0')
_______________Net 2_______________
V0902 15:57:24.068000 140183681587008 torch/_inductor/scheduler.py:1820] [0/2] [__fusion] ===== attempting fusion (1/10): 2 nodes =====
V0902 15:57:24.069000 140183681587008 torch/_inductor/scheduler.py:627] [0/2] [__fusion] cannot fuse buf0 with buf1: numel/rnumel mismatch (reduce) (123, 1), (((s0**3 + 122)//123), 123)
V0902 15:57:24.070000 140183681587008 torch/_inductor/scheduler.py:2133] [0/2] [__fusion] found 0 possible fusions
V0902 15:57:24.070000 140183681587008 torch/_inductor/scheduler.py:1827] [0/2] [__fusion] completed fusion round (1/10): fused 2 nodes into 2 nodes
V0902 15:57:24.070000 140183681587008 torch/_inductor/scheduler.py:1827] [0/2] [__fusion] 
V0902 15:57:24.070000 140183681587008 torch/_inductor/scheduler.py:1834] [0/2] [__fusion] ===== fusion complete (1 iterations) =====
V0902 15:57:24.479000 140183681587008 torch/_inductor/scheduler.py:1820] [0/2] [__fusion] ===== attempting fusion (1/10): 1 nodes =====
V0902 15:57:24.479000 140183681587008 torch/_inductor/scheduler.py:2133] [0/2] [__fusion] found 0 possible fusions
V0902 15:57:24.480000 140183681587008 torch/_inductor/scheduler.py:1827] [0/2] [__fusion] completed fusion round (1/10): fused 1 nodes into 1 nodes
V0902 15:57:24.480000 140183681587008 torch/_inductor/scheduler.py:1827] [0/2] [__fusion] 
V0902 15:57:24.480000 140183681587008 torch/_inductor/scheduler.py:1834] [0/2] [__fusion] ===== fusion complete (1 iterations) =====
Time: 1076.0277099609375 ms
tensor(0.1223, device='cuda:0')
_______________Compiled Graphed Iteration 0_______________
Time: 0.20978688031435014 ms
tensor(0.3670, device='cuda:0')
_______________Compiled Graphed Iteration 1_______________
Time: 0.20953951954841613 ms
tensor(0.3670, device='cuda:0')
_______________Compiled Graphed Iteration 2_______________
Time: 0.2100307197868824 ms
tensor(0.3670, device='cuda:0')
```

I expect that only Net 0 is compiled since Net 1 and Net 2 are different instances of A with the same attributes.
Note, I request `recompiles=True` in logging, but it does not show anything, however from fusion attempts and the timing I can tell it tries to recompile. 

If I request `guards=True` in `set_logs`, I see
```
[__guards] | | +- ID_MATCH: ___check_obj_id(L['net'], 140280864840976)
```
which suggests that the compiled code is assumed valid only for exactly the same object of type inheriting nn.Module.

If I let 
```
Mod = object
```
Net 1 and Net 2 iterations correctly use the compiled code:

```
_______________Net 0_______________
V0902 16:06:52.766000 140478342887232 torch/_inductor/scheduler.py:1820] [0/0] [__fusion] ===== attempting fusion (1/10): 2 nodes =====
V0902 16:06:52.767000 140478342887232 torch/_inductor/scheduler.py:627] [0/0] [__fusion] cannot fuse buf0 with buf1: numel/rnumel mismatch (reduce) (123, 1), (((s0**3 + 122)//123), 123)
V0902 16:06:52.768000 140478342887232 torch/_inductor/scheduler.py:2133] [0/0] [__fusion] found 0 possible fusions
V0902 16:06:52.769000 140478342887232 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] completed fusion round (1/10): fused 2 nodes into 2 nodes
V0902 16:06:52.769000 140478342887232 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] 
V0902 16:06:52.769000 140478342887232 torch/_inductor/scheduler.py:1834] [0/0] [__fusion] ===== fusion complete (1 iterations) =====
V0902 16:06:54.450000 140478342887232 torch/_inductor/scheduler.py:1820] [0/0] [__fusion] ===== attempting fusion (1/10): 1 nodes =====
V0902 16:06:54.451000 140478342887232 torch/_inductor/scheduler.py:2133] [0/0] [__fusion] found 0 possible fusions
V0902 16:06:54.451000 140478342887232 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] completed fusion round (1/10): fused 1 nodes into 1 nodes
V0902 16:06:54.451000 140478342887232 torch/_inductor/scheduler.py:1827] [0/0] [__fusion] 
V0902 16:06:54.451000 140478342887232 torch/_inductor/scheduler.py:1834] [0/0] [__fusion] ===== fusion complete (1 iterations) =====
Time: 4958.19677734375 ms
tensor(0.1222, device='cuda:0')
_______________Net 1_______________
Time: 1.074720025062561 ms
tensor(0.1222, device='cuda:0')
_______________Net 2_______________
Time: 0.5718399882316589 ms
tensor(0.1222, device='cuda:0')
_______________Compiled Graphed Iteration 0_______________
Time: 0.20977663949131967 ms
tensor(0.3666, device='cuda:0')
_______________Compiled Graphed Iteration 1_______________
Time: 0.20968575939536094 ms
tensor(0.3666, device='cuda:0')
_______________Compiled Graphed Iteration 2_______________
Time: 0.2093523195385933 ms
tensor(0.3666, device='cuda:0')
```

And finally I found a crazy workaround to actually allow nn.Module in arguments, by duplicating the class thro re-importing it:

```python 
from importlib import reload
def get_Mod():
    mod = reload(torch.nn.modules.module)
    return mod.Module
Mod = get_Mod()
```
This gives the same output as above for the case `Mod=object`. So it seems that dynamo has some special logic particularly for nn.Module?

My real use case is compiling a block which is a list of layers inheriting from nn.Module. he network contains multiple blocks of the same kind but differing in values of some parameter / buffer tensors.

There is a relevant case for the recurrent neural network as well:
https://discuss.pytorch.org/t/how-to-compile-a-rnn-loop-once/191455?u=alexander_shekhovts



### Versions

Collecting environment information...
PyTorch version: 2.4.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (GCC) 13.2.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.11.5 (main, Oct  2 2023, 09:22:39) [GCC 13.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-1062-nvidia-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA DGX Display
GPU 4: NVIDIA A100-SXM4-40GB

Nvidia driver version: 535.183.01
cuDNN version: Could not collect
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:                        43 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               128
On-line CPU(s) list:                  0-127
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7742 64-Core Processor
CPU family:                           23
Model:                                49
Thread(s) per core:                   2
Core(s) per socket:                   64
Socket(s):                            1
Stepping:                             0
Frequency boost:                      enabled
CPU max MHz:                          2250,0000
CPU min MHz:                          1500,0000
BogoMIPS:                             4491.73
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es
Virtualization:                       AMD-V
L1d cache:                            2 MiB (64 instances)
L1i cache:                            2 MiB (64 instances)
L2 cache:                             32 MiB (64 instances)
L3 cache:                             256 MiB (16 instances)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-127
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow:   Mitigation; safe RET
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==2.1.0
[pip3] torch==2.4.0
[pip3] torchvision==0.19.0
[pip3] triton==3.0.0
[conda] Could not collect


cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec