Repository: pytorch
Issue: Fails to compile when using DistributedDataParallel: AttributeError: 'int' object has no attribute 'meta'
Labels: needs reproduction,oncall: distributed,triaged,module: ddp,oncall: pt2,pt2d-triage-nov2024

### 🐛 Describe the bug

see the log outputs here: https://gist.github.com/alita-moore/a23f57d2adfaa885185308e922aef44a

To reproduce this run the following code:

```
import os
import torch
import torch.distributed as dist
from timm.models.swin_transformer_v2 import SwinTransformerV2, WindowAttention

def main():
    # Initialize the distributed environment
    dist.init_process_group(backend='nccl')  # Use NCCL backend
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    # Set the device
    if torch.cuda.is_available():
        device = torch.device(f"cuda:{rank}")
        torch.cuda.set_device(device)
    else:
        device = torch.device("cpu")
    
    print(f"Running on rank {rank} out of {world_size} processes, device: {device}")

    # Compile WindowAttention.forward
    WindowAttention.forward = torch.compile(WindowAttention.forward)

    # Create the model
    model = SwinTransformerV2(
        img_size=(128, 128),
        depths=(2, 2, 18, 2),
        window_size=8,
        patch_size=4,
        embed_dim=128,
        num_heads=(4, 8, 16, 32),
        num_classes=0,
        in_chans=3
    ).to(device)

    # Wrap the model with DistributedDataParallel
    ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank] if torch.cuda.is_available() else None)

    # Create a sample input tensor
    im = torch.empty(1, 3, 128, 128).to(device)

    # Forward pass
    output = ddp_model(im)
    print(f"Rank {rank}: Output shape: {output.shape}")

    # Clean up
    dist.destroy_process_group()

if __name__ == "__main__":
    main()
```

Then call it with `torchrun --nproc_per_node=1 /path/to/file.py`

### Versions

Collecting environment information...
PyTorch version: 2.3.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 12 (bookworm) (x86_64)
GCC version: (Debian 12.2.0-14) 12.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.36

Python version: 3.10.13 (main, Mar 12 2024, 12:22:40) [GCC 12.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-1013-aws-x86_64-with-glibc2.36
Is CUDA available: True
CUDA runtime version: 12.3.107
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L4
Nvidia driver version: 555.58.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.3.0
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:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               4
On-line CPU(s) list:                  0-3
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7R13 Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   2
Socket(s):                            1
Stepping:                             1
BogoMIPS:                             5299.99
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 tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save vaes vpclmulqdq rdpid
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            64 KiB (2 instances)
L1i cache:                            64 KiB (2 instances)
L2 cache:                             1 MiB (2 instances)
L3 cache:                             8 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-3
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:               Not affected
Vulnerability Spec rstack overflow:   Vulnerable: Safe RET, no microcode
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; 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] mypy==1.11.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] torch==2.3.1
[pip3] torchaudio==2.3.1
[pip3] torchtext==0.18.0
[pip3] torchvision==0.18.1
[pip3] triton==2.3.1
[conda] Could not collect

----


NOTE: I'm using timm version ^1.0.8 (https://github.com/huggingface/pytorch-image-models)

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @ezyang @chauhang @penguinwu